Repository: fsspec/filesystem_spec
Pull request #1162 (2023-01-20 18:34:50)

Explicit writing for ZipFileSystem
Note: this changes the default in `ls` to `detail=True`, as it is in the spec and should be for all FS implementations.


Patch:
diff --git a/fsspec/archive.py b/fsspec/archive.py
index 79aa411..dc5c149 100644
--- a/fsspec/archive.py
+++ b/fsspec/archive.py
@@ -46,7 +46,7 @@ class AbstractArchiveFileSystem(AbstractFileSystem):
         else:
             raise FileNotFoundError(path)
 
-    def ls(self, path, detail=False, **kwargs):
+    def ls(self, path, detail=True, **kwargs):
         self._get_dirs()
         paths = {}
         for p, f in self.dir_cache.items():
@@ -66,8 +66,8 @@ class AbstractArchiveFileSystem(AbstractFileSystem):
                 if ppath not in paths:
                     out = {"name": ppath + "/", "size": 0, "type": "directory"}
                     paths[ppath] = out
-        out = list(paths.values())
+        out = sorted(paths.values(), key=lambda _: _["name"])
         if detail:
             return out
         else:
-            return list(sorted(f["name"] for f in out))
+            return [f["name"] for f in out]
diff --git a/fsspec/implementations/zip.py b/fsspec/implementations/zip.py
index 71257e0..cb547f0 100644
--- a/fsspec/implementations/zip.py
+++ b/fsspec/implementations/zip.py
@@ -1,15 +1,13 @@
 from __future__ import absolute_import, division, print_function
 
-import datetime
 import zipfile
 
 import fsspec
 from fsspec.archive import AbstractArchiveFileSystem
-from fsspec.utils import DEFAULT_BLOCK_SIZE
 
 
 class ZipFileSystem(AbstractArchiveFileSystem):
-    """Read contents of ZIP archive as a file-system
+    """Read/Write contents of ZIP archive as a file-system
 
     Keeps file object open while instance lives.
 
@@ -26,7 +24,9 @@ class ZipFileSystem(AbstractArchiveFileSystem):
         mode="r",
         target_protocol=None,
         target_options=None,
-        block_size=DEFAULT_BLOCK_SIZE,
+        compression=zipfile.ZIP_STORED,
+        allowZip64=True,
+        compresslevel=None,
         **kwargs,
     ):
         """
@@ -36,23 +36,33 @@ class ZipFileSystem(AbstractArchiveFileSystem):
             Contains ZIP, and must exist. If a str, will fetch file using
             :meth:`~fsspec.open_files`, which must return one file exactly.
         mode: str
-            Currently, only 'r' accepted
+            Accept: "r", "w", "a"
         target_protocol: str (optional)
             If ``fo`` is a string, this value can be used to override the
             FS protocol inferred from a URL
         target_options: dict (optional)
             Kwargs passed when instantiating the target FS, if ``fo`` is
             a string.
+        compression, allowZip64, compresslevel: passed to ZipFile
+            Only relevant when creating a ZIP
         """
         super().__init__(self, **kwargs)
+        if mode not in set("rwa"):
+            raise ValueError(f"mode '{mode}' no understood")
+        self.mode = mode
         if isinstance(fo, str):
             fo = fsspec.open(
                 fo, mode=mode + "b", protocol=target_protocol, **(target_options or {})
             )
         self.of = fo
         self.fo = fo.__enter__()  # the whole instance is a context
-        self.zip = zipfile.ZipFile(self.fo, mode=mode)
-        self.block_size = block_size
+        self.zip = zipfile.ZipFile(
+            self.fo,
+            mode=mode,
+            compression=compression,
+            allowZip64=allowZip64,
+            compresslevel=compresslevel,
+        )
         self.dir_cache = None
 
     @classmethod
@@ -63,13 +73,16 @@ class ZipFileSystem(AbstractArchiveFileSystem):
     def __del__(self):
         if hasattr(self, "zip"):
             self.close()
+        del self.zip
 
     def close(self):
-        "Commits any write changes to the file. Done on ``del`` too."
+        """Commits any write changes to the file. Done on ``del`` too."""
         self.zip.close()
 
     def _get_dirs(self):
-        if self.dir_cache is None:
+        if self.dir_cache is None or self.mode in set("wa"):
+            # when writing, dir_cache is always in the ZipFile's attributes,
+            # not read from the file.
             files = self.zip.infolist()
             self.dir_cache = {
                 dirname + "/": {"name": dirname + "/", "size": 0, "type": "directory"}
@@ -88,10 +101,7 @@ class ZipFileSystem(AbstractArchiveFileSystem):
 
     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)
+        self.zip.writestr(path, value, **kwargs)
 
     def _open(
         self,
@@ -103,6 +113,12 @@ class ZipFileSystem(AbstractArchiveFileSystem):
         **kwargs,
     ):
         path = self._strip_protocol(path)
+        if "r" in mode and self.mode in set("wa"):
+            if self.exists(path):
+                raise IOError("ZipFS can only be open for reading or writing, not both")
+            raise FileNotFoundError(path)
+        if "r" in self.mode and "w" in mode:
+            raise IOError("ZipFS can only be open for reading or writing, not both")
         out = self.zip.open(path, mode.strip("b"))
         if "r" in mode:
             info = self.info(path)


Test patch:
diff --git a/fsspec/implementations/tests/test_archive.py b/fsspec/implementations/tests/test_archive.py
index 2f7fe70..312e4a2 100644
--- a/fsspec/implementations/tests/test_archive.py
+++ b/fsspec/implementations/tests/test_archive.py
@@ -276,13 +276,13 @@ class TestAnyArchive:
         with scenario.provider(archive_data) as archive:
             fs = fsspec.filesystem(scenario.protocol, fo=archive)
 
-            assert fs.ls("") == ["a", "b", "deeply/"]
+            assert fs.ls("", detail=False) == ["a", "b", "deeply/"]
             assert fs.ls("/") == fs.ls("")
 
-            assert fs.ls("deeply") == ["deeply/nested/"]
+            assert fs.ls("deeply", detail=False) == ["deeply/nested/"]
             assert fs.ls("deeply/") == fs.ls("deeply")
 
-            assert fs.ls("deeply/nested") == ["deeply/nested/path"]
+            assert fs.ls("deeply/nested", detail=False) == ["deeply/nested/path"]
             assert fs.ls("deeply/nested/") == fs.ls("deeply/nested")
 
     def test_find(self, scenario: ArchiveTestScenario):
diff --git a/fsspec/implementations/tests/test_zip.py b/fsspec/implementations/tests/test_zip.py
index 8fe0d48..93c758f 100644
--- a/fsspec/implementations/tests/test_zip.py
+++ b/fsspec/implementations/tests/test_zip.py
@@ -1,5 +1,7 @@
 import collections.abc
 
+import pytest
+
 import fsspec
 from fsspec.implementations.tests.test_archive import archive_data, tempzip
 
@@ -57,3 +59,27 @@ def test_write_seek(m):
     with m.open("afile.zip", "rb") as f:
         fs = fsspec.filesystem("zip", fo=f)
         assert fs.cat("another") == b"hi"
+
+
+def test_rw(m):
+    # extra arg to zip means "create archive"
+    with fsspec.open(
+        "zip://afile::memory://out.zip", mode="wb", zip={"mode": "w"}
+    ) as f:
+        f.write(b"data")
+
+    with fsspec.open("zip://afile::memory://out.zip", mode="rb") as f:
+        assert f.read() == b"data"
+
+
+def test_mapper(m):
+    # extra arg to zip means "create archive"
+    mapper = fsspec.get_mapper("zip::memory://out.zip", zip={"mode": "w"})
+    with pytest.raises(KeyError):
+        mapper["a"]
+
+    mapper["a"] = b"data"
+    with pytest.raises(OSError):
+        # fails because this is write mode and we cannot also read
+        mapper["a"]
+    assert "a" in mapper  # but be can list