Repository: fsspec/filesystem_spec
Pull request #1017 (2022-08-08 17:25:20)

Allow writable ZipFileSystem

martindurant: It appears to be necessary for the calling code to call `fs.zip.close()` to get the output to finalize, I am not sure why.

Patch:
diff --git a/fsspec/implementations/zip.py b/fsspec/implementations/zip.py
index 2ede0dc..30c18cb 100644
--- a/fsspec/implementations/zip.py
+++ b/fsspec/implementations/zip.py
@@ -1,5 +1,6 @@
 from __future__ import absolute_import, division, print_function
 
+import datetime
 import zipfile
 
 from fsspec import open_files
@@ -44,10 +45,10 @@ class ZipFileSystem(AbstractArchiveFileSystem):
             a string.
         """
         super().__init__(self, **kwargs)
-        if mode != "r":
-            raise ValueError("Only read from zip files accepted")
         if isinstance(fo, str):
-            files = open_files(fo, protocol=target_protocol, **(target_options or {}))
+            files = open_files(
+                fo, mode=mode + "b", protocol=target_protocol, **(target_options or {})
+            )
             if len(files) != 1:
                 raise ValueError(
                     'Path "{}" did not resolve to exactly'
@@ -55,7 +56,7 @@ class ZipFileSystem(AbstractArchiveFileSystem):
                 )
             fo = files[0]
         self.fo = fo.__enter__()  # the whole instance is a context
-        self.zip = zipfile.ZipFile(self.fo)
+        self.zip = zipfile.ZipFile(self.fo, mode=mode)
         self.block_size = block_size
         self.dir_cache = None
 
@@ -64,6 +65,9 @@ class ZipFileSystem(AbstractArchiveFileSystem):
         # zip file paths are always relative to the archive root
         return super()._strip_protocol(path).lstrip("/")
 
+    def __del__(self):
+        self.zip.close()
+
     def _get_dirs(self):
         if self.dir_cache is None:
             files = self.zip.infolist()
@@ -82,6 +86,13 @@ class ZipFileSystem(AbstractArchiveFileSystem):
                 )
                 self.dir_cache[f["name"]] = f
 
+    def pipe_file(self, path, value, **kwargs):
+        # override upstream, because we know the exact file size in this case
+        info = zipfile.ZipInfo(path, datetime.datetime.now().timetuple())
+        info.file_size = len(value)
+        with self.zip.open(path, "w") as f:
+            f.write(value)
+
     def _open(
         self,
         path,
@@ -92,10 +103,9 @@ class ZipFileSystem(AbstractArchiveFileSystem):
         **kwargs,
     ):
         path = self._strip_protocol(path)
-        if mode != "rb":
-            raise NotImplementedError
-        info = self.info(path)
-        out = self.zip.open(path, "r")
-        out.size = info["size"]
-        out.name = info["name"]
+        out = self.zip.open(path, mode.strip("b"))
+        if "r" in mode:
+            info = self.info(path)
+            out.size = info["size"]
+            out.name = info["name"]
         return out
diff --git a/setup.py b/setup.py
index e147293..3e4e55a 100644
--- a/setup.py
+++ b/setup.py
@@ -50,7 +50,7 @@ setup(
         "gs": ["gcsfs"],
         "hdfs": ["pyarrow >= 1"],
         "arrow": ["pyarrow >= 1"],
-        "http": ["requests", "aiohttp"],
+        "http": ["requests", "aiohttp !=4.0.0a0, !=4.0.0a1"],
         "sftp": ["paramiko"],
         "s3": ["s3fs"],
         "oci": ["ocifs"],


Test patch:
diff --git a/fsspec/implementations/tests/test_zip.py b/fsspec/implementations/tests/test_zip.py
index 6be4013..8fe0d48 100644
--- a/fsspec/implementations/tests/test_zip.py
+++ b/fsspec/implementations/tests/test_zip.py
@@ -46,3 +46,14 @@ def test_root_info():
         fs = fsspec.filesystem("zip", fo=z)
         assert fs.info("/") == {"name": "/", "type": "directory", "size": 0}
         assert fs.info("") == {"name": "/", "type": "directory", "size": 0}
+
+
+def test_write_seek(m):
+    with m.open("afile.zip", "wb") as f:
+        fs = fsspec.filesystem("zip", fo=f, mode="w")
+        fs.pipe("another", b"hi")
+        fs.zip.close()
+
+    with m.open("afile.zip", "rb") as f:
+        fs = fsspec.filesystem("zip", fo=f)
+        assert fs.cat("another") == b"hi"