Merge pull request #25 from findmyway/rever_part_of_24

add back record_type in read

Project.toml

@@ -1,7 +1,7 @@
 name = "TFRecord"
 uuid = "841416d8-1a6a-485a-b0fc-1328d0f53d5e"
 authors = ["Jun Tian <[email protected]> and contributors"]
-version = "0.4.0"
+version = "0.5.0"
 
 [deps]
 BufferedStreams = "e1450e63-4bb3-523b-b2a4-4ffa8c0fd77d"

src/core.jl

@@ -51,6 +51,7 @@ Read tensorflow records from file(s).
 - `compression=nothing`. No compression by default. Optional values are `:zlib` and `:gzip`.
 - `bufsize=10*1024*1024`. Set the buffer size of internal `BufferedOutputStream`. The default value is `10M`. Suggested value is between `1M`~`100M`.
 - `channel_size=1000`. The number of pre-fetched elements.
+- `record_type=Example`, see https://github.com/JuliaReinforcementLearning/TFRecord.jl/pull/11
 
 !!! note
 
@@ -61,15 +62,17 @@ function read(
     compression = nothing,
     bufsize = 10 * 1024 * 1024,
     channel_size = 1_000,
+    record_type = Example,
+
 )
     file_itr(file::AbstractString) = [file]
     file_itr(files) = files
-    Channel{Example}(channel_size) do ch
+    Channel{record_type}(channel_size) do ch
         @threads for file_name in file_itr(files)
             open(decompressor_stream(compression), file_name, "r") do io
                 buffered_io = BufferedInputStream(io, bufsize)
                 while !eof(buffered_io)
-                    instance = readproto(IOBuffer(read_record(buffered_io)), Example())
+                    instance = readproto(IOBuffer(read_record(buffered_io)), record_type())
                     put!(ch, instance)
                 end
             end