From 6f9546a7397cabc3692620d298c49a0abc15cb2f Mon Sep 17 00:00:00 2001
From: nhz2 <nhz2@cornell.edu>
Date: Wed, 12 Jul 2023 13:58:03 -0400
Subject: [PATCH] change zip reader and writer

---
 src/loading.jl |  2 +-
 src/readers.jl | 11 +++++------
 src/saving.jl  | 22 ++++++++++++++++------
 src/writers.jl | 23 +++++++----------------
 4 files changed, 29 insertions(+), 29 deletions(-)

diff --git a/src/loading.jl b/src/loading.jl
index d89d8ef..9c583f4 100644
--- a/src/loading.jl
+++ b/src/loading.jl
@@ -4,7 +4,7 @@ function load_dir(dirpath::AbstractString)::ZGroup
     reader = if isdir(dirpath)
         DirectoryReader(dirpath)
     elseif isfile(dirpath)
-        BufferedZipReader(dirpath)
+        ZarrZipReader(read(dirpath))
     else
         throw(ArgumentError("loading directory $(repr(dirpath)): No such file or directory"))
     end
diff --git a/src/readers.jl b/src/readers.jl
index 0308231..f9dd631 100644
--- a/src/readers.jl
+++ b/src/readers.jl
@@ -46,18 +46,17 @@ function read_key_idx(d::DirectoryReader, idx::Int)::Vector{UInt8}
 end
 
 
-struct BufferedZipReader <: AbstractReader
+struct ZarrZipReader <: AbstractReader
     zipfile::ZipBufferReader{Vector{UInt8}}
-    function BufferedZipReader(path)
-        @argcheck isfile(path)
-        new(ZipBufferReader(read(path)))
+    function ZarrZipReader(buffer::Vector{UInt8})
+        new(ZipBufferReader(buffer))
     end
 end
 
-function key_names(d::BufferedZipReader)::Vector{String}
+function key_names(d::ZarrZipReader)::Vector{String}
     return zip_names(d.zipfile)
 end
 
-function read_key_idx(d::BufferedZipReader, idx::Int)::Vector{UInt8}
+function read_key_idx(d::ZarrZipReader, idx::Int)::Vector{UInt8}
     zip_readentry(d.zipfile, idx)
 end
\ No newline at end of file
diff --git a/src/saving.jl b/src/saving.jl
index 7733307..23f4b4c 100644
--- a/src/saving.jl
+++ b/src/saving.jl
@@ -1,18 +1,28 @@
 
+
 """
 If dirpath ends in .zip, save to a zip file, otherwise save to a directory.
+
+Note this will delete pre existing data at dirpath
 """
 function save_dir(dirpath::AbstractString, z::ZGroup)
-    writer = if endswith(dirpath, ".zip")
-        BufferedZipWriter(dirpath)
+    if endswith(dirpath, ".zip")
+        @argcheck !isdir(dirpath)
+        mkpath(dirname(dirpath))
+        open(dirpath; write=true) do io
+            writer = ZarrZipWriter(io)
+            try
+                save_dir(writer, z)
+            finally
+                closewriter(writer)
+            end
+        end
     else
-        DirectoryWriter(dirpath)
+        save_dir(DirectoryWriter(dirpath), z)
     end
-    save_dir(writer, z)
-    closewriter(writer)
+    nothing
 end
 
-
 """
 Note this will delete pre existing data at dirpath
 """
diff --git a/src/writers.jl b/src/writers.jl
index cd3571e..8525280 100644
--- a/src/writers.jl
+++ b/src/writers.jl
@@ -40,28 +40,19 @@ end
 Write to an in memory zipfile, that gets saved to disk on close.
 This writer will overwrite any existing file at `path`
 """
-struct BufferedZipWriter <: AbstractWriter
-    zipfile::ZipWriter{IOBuffer}
-    path::String
-    iobuffer::IOBuffer
-    function BufferedZipWriter(path)
-        @argcheck !isdir(path)
-        @argcheck !isdirpath(path)
-        mkpath(dirname(path))
-        iobuffer = IOBuffer()
-        zipfile = ZipWriter(iobuffer)
-        new(zipfile, abspath(path), iobuffer)
+struct ZarrZipWriter{IO_TYPE<:IO} <: AbstractWriter
+    zipfile::ZipWriter{IO_TYPE}
+    function ZarrZipWriter(io)
+        new{typeof(io)}(ZipWriter(io))
     end
 end
 
-function write_key(d::BufferedZipWriter, key::AbstractString, data)::Nothing
+function write_key(d::ZarrZipWriter, key::AbstractString, data)::Nothing
     zip_writefile(d.zipfile, key, data)
     nothing
 end
 
-function closewriter(d::BufferedZipWriter)
+function closewriter(d::ZarrZipWriter)
     close(d.zipfile)
-    open(d.path, "w") do f
-        write(f, take!(d.iobuffer))
-    end
+    nothing
 end
\ No newline at end of file