feat: ZipStore accepts open binary file-like objects

changes/4187.feature.md

@@ -0,0 +1,4 @@
+`ZipStore` now accepts an open binary file-like object in place of a path, enabling
+zip archives on remote storage (e.g. a file opened with `fsspec` or an
+`obstore.ReadableFile`). Operations that require a filesystem location
+(`clear`, `move`) raise `NotImplementedError` for file-object-backed stores.

docs/user-guide/storage.md

@@ -124,6 +124,19 @@ array = zarr.create_array(store=store, shape=(2,), dtype='float64')
 print(array)
 ```
 
+In place of a path, `ZipStore` also accepts an open binary file object (for
+example a file opened with `fsspec`, or an `obstore` reader), enabling zip
+archives on remote storage. The file must stay open for as long as the store
+is in use:
+
+```python exec="true" session="storage" source="above" result="ansi"
+store.close()
+f = open('data.zip', mode='rb')  # must stay open while the store is used
+array = zarr.open_array(store=zarr.storage.ZipStore(f), mode='r')
+print(array[:])
+f.close()
+```
+
 ### Remote Store
 
 The [`zarr.storage.FsspecStore`][] stores the contents of a Zarr hierarchy following the same

src/zarr/storage/_zip.py

@@ -1,12 +1,13 @@
 from __future__ import annotations
 
+import io
 import os
 import shutil
 import threading
 import time
 import zipfile
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Literal
+from typing import IO, TYPE_CHECKING, Any, Literal
 
 from zarr.abc.store import (
     ByteRequest,
@@ -23,14 +24,67 @@
 ZipStoreAccessModeLiteral = Literal["r", "w", "a"]
 
 
+class _RawReaderAdapter(io.RawIOBase):
+    """
+    Adapt a minimal seekable reader to the `io` interface `zipfile` needs.
+
+    Some file-like objects (e.g. `obstore.ReadableFile`) implement
+    `read`/`seek`/`tell` but are not `io.IOBase` instances, and their
+    `read` may return a buffer-protocol object rather than `bytes`.
+    Wrapping in this adapter plus `io.BufferedReader` yields real `bytes`.
+
+    Reads are clamped to the bytes remaining before EOF: some readers
+    (obstore < 0.6) raise on short reads rather than returning fewer bytes.
+    The size is cached, which is safe because the adapter is only used for
+    read-only access.
+    """
+
+    def __init__(self, fileobj: IO[bytes]) -> None:
+        self._fileobj = fileobj
+        self._size: int | None = None
+
+    def _get_size(self) -> int:
+        if self._size is None:
+            pos = self._fileobj.tell()
+            self._size = self._fileobj.seek(0, os.SEEK_END)
+            self._fileobj.seek(pos)
+        return self._size
+
+    def readable(self) -> bool:
+        return True
+
+    def seekable(self) -> bool:
+        return True
+
+    def seek(self, pos: int, whence: int = 0) -> int:
+        return self._fileobj.seek(pos, whence)
+
+    def tell(self) -> int:
+        return self._fileobj.tell()
+
+    def readinto(self, b: Any) -> int:
+        n_requested = min(len(b), self._get_size() - self._fileobj.tell())
+        if n_requested <= 0:
+            return 0
+        data = self._fileobj.read(n_requested)
+        n = len(data)
+        b[:n] = memoryview(data)
+        return n
+
+
 class ZipStore(Store):
     """
     Store using a ZIP file.
 
     Parameters
     ----------
-    path : str
-        Location of file.
+    path : str, Path, or IO[bytes]
+        Location of file, or an open binary file object. A file object must
+        support `read`, `seek`, and `tell`; objects that are not `io.IOBase`
+        instances (e.g. an `obstore` reader) are adapted automatically but
+        can only be used for reading (`mode="r"`). The file object must stay
+        open for the lifetime of the store, and operations that require a
+        filesystem location (`clear`, `move`, pickling) are not supported.
     mode : str, optional
         One of 'r' to read an existing file, 'w' to truncate and write a new
         file, 'a' to append to an existing file, or 'x' to exclusively create
@@ -58,16 +112,17 @@ class ZipStore(Store):
     supports_deletes: bool = False
     supports_listing: bool = True
 
-    path: Path
+    path: Path | None
     compression: int
     allowZip64: bool
 
     _zf: zipfile.ZipFile
     _lock: threading.RLock
+    _fileobj: IO[bytes] | None
 
     def __init__(
         self,
-        path: Path | str,
+        path: Path | str | IO[bytes],
         *,
         mode: ZipStoreAccessModeLiteral = "r",
         read_only: bool | None = None,
@@ -81,8 +136,28 @@ def __init__(
 
         if isinstance(path, str):
             path = Path(path)
-        assert isinstance(path, Path)
-        self.path = path  # root?
+        if isinstance(path, Path):
+            self.path = path  # root?
+            self._fileobj = None
+        else:
+            self.path = None
+            if not isinstance(path, io.IOBase):
+                if not all(
+                    callable(getattr(path, attr, None)) for attr in ("read", "seek", "tell")
+                ):
+                    raise TypeError(
+                        f"expected a path or an open binary file object supporting "
+                        f"read/seek/tell, got {type(path).__name__}"
+                    )
+                if mode != "r":
+                    raise TypeError(
+                        f"a file object that is not an io.IOBase instance can only be "
+                        f"opened for reading (mode='r', got mode={mode!r})"
+                    )
+                # e.g. an obstore ReadableFile: readable and seekable, but
+                # not an io object and reads may not return bytes
+                path = io.BufferedReader(_RawReaderAdapter(path))
+            self._fileobj = path
 
         self._zmode = mode
         self.compression = compression
@@ -95,7 +170,7 @@ def _sync_open(self) -> None:
         self._lock = threading.RLock()
 
         self._zf = zipfile.ZipFile(
-            self.path,
+            self.path if self.path is not None else self._fileobj,  # type: ignore[arg-type]
             mode=self._zmode,
             compression=self.compression,
             allowZip64=self.allowZip64,
@@ -107,6 +182,13 @@ async def _open(self) -> None:
         self._sync_open()
 
     def __getstate__(self) -> dict[str, Any]:
+        if self.path is None:
+            # A path-backed store pickles its path and reopens the file on
+            # unpickling; an open file object cannot be serialized that way.
+            raise TypeError(
+                "cannot pickle a ZipStore backed by a file-like object; "
+                "construct the store from a path instead"
+            )
         # We need a copy to not modify the state of the original store
         state = self.__dict__.copy()
         for attr in ["_zf", "_lock"]:
@@ -130,20 +212,30 @@ async def clear(self) -> None:
         # docstring inherited
         with self._lock:
             self._check_writable()
+            if self.path is None:
+                raise NotImplementedError(
+                    "clear() is not supported for a ZipStore backed by a file-like object"
+                )
             self._zf.close()
             os.remove(self.path)
             self._zf = zipfile.ZipFile(
                 self.path, mode="w", compression=self.compression, allowZip64=self.allowZip64
             )
 
     def __str__(self) -> str:
+        if self.path is None:
+            return f"zip://{self._fileobj!r}"
         return f"zip://{self.path}"
 
     def __repr__(self) -> str:
         return f"ZipStore('{self}')"
 
     def __eq__(self, other: object) -> bool:
-        return isinstance(other, type(self)) and self.path == other.path
+        return (
+            isinstance(other, type(self))
+            and self.path == other.path
+            and self._fileobj is other._fileobj
+        )
 
     def _get(
         self,
@@ -297,6 +389,10 @@ async def move(self, path: Path | str) -> None:
         """
         Move the store to another path.
         """
+        if self.path is None:
+            raise NotImplementedError(
+                "move() is not supported for a ZipStore backed by a file-like object"
+            )
         if isinstance(path, str):
             path = Path(path)
         self.close()

tests/test_store/test_zip.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
+import io
 import os
+import pickle
 import shutil
 import tempfile
 import zipfile
@@ -188,6 +190,163 @@ async def test_move(self, tmp_path: Path) -> None:
         assert np.array_equal(array[...], np.arange(10))
 
 
+class TestZipStoreFileObj:
+    """ZipStore backed by an open binary file-like object instead of a path."""
+
+    @pytest.fixture
+    def zip_bytes(self, tmp_path: Path) -> bytes:
+        path = tmp_path / "data.zip"
+        store = ZipStore(path, mode="w")
+        zarr.create_array(store, data=np.arange(10), chunks=(5,))
+        store.close()
+        return path.read_bytes()
+
+    def test_read_from_fileobj(self, zip_bytes: bytes) -> None:
+        # an existing archive can be read through any seekable binary reader
+        store = ZipStore(io.BytesIO(zip_bytes), mode="r")
+        array = zarr.open_array(store, mode="r")
+        assert np.array_equal(array[...], np.arange(10))
+        assert store.path is None
+
+    def test_write_to_fileobj(self) -> None:
+        # a writable file object receives the archive; the bytes it holds
+        # after close() are a complete, reopenable zip
+        buffer = io.BytesIO()
+        store = ZipStore(buffer, mode="w", read_only=False)
+        zarr.create_array(store, data=np.arange(4))
+        store.close()
+
+        roundtrip = ZipStore(io.BytesIO(buffer.getvalue()), mode="r")
+        array = zarr.open_array(roundtrip, mode="r")
+        assert np.array_equal(array[...], np.arange(4))
+
+    async def test_clear_unsupported(self, zip_bytes: bytes) -> None:
+        # clear() requires a filesystem location, so it raises a clear error
+        # for file-object-backed stores
+        store = ZipStore(io.BytesIO(zip_bytes), mode="a", read_only=False)
+        store._sync_open()
+        with pytest.raises(NotImplementedError, match="clear.*file-like"):
+            await store.clear()
+
+    async def test_move_unsupported(self, zip_bytes: bytes) -> None:
+        # move() requires a filesystem location, so it raises a clear error
+        # for file-object-backed stores
+        store = ZipStore(io.BytesIO(zip_bytes), mode="a", read_only=False)
+        store._sync_open()
+        with pytest.raises(NotImplementedError, match="move.*file-like"):
+            await store.move("elsewhere.zip")
+
+    def test_invalid_file_object_rejected(self) -> None:
+        # objects without read/seek/tell are rejected at construction, not
+        # deep inside zipfile
+        with pytest.raises(TypeError, match="read/seek/tell"):
+            ZipStore(42, mode="r")  # type: ignore[arg-type]
+
+    @pytest.mark.parametrize("mode", ["w", "a", "x"])
+    def test_non_iobase_reader_write_modes_rejected(self, zip_bytes: bytes, mode: str) -> None:
+        # readers that are not io.IOBase instances are adapted for reading
+        # only; write modes are rejected at construction with a clear error
+        class MinimalReader:
+            def __init__(self, data: bytes) -> None:
+                self._buffer = io.BytesIO(data)
+
+            def read(self, size: int, /) -> bytes:
+                return self._buffer.read(size)
+
+            def seek(self, pos: int, whence: int = 0, /) -> int:
+                return self._buffer.seek(pos, whence)
+
+            def tell(self) -> int:
+                return self._buffer.tell()
+
+        with pytest.raises(TypeError, match="opened for reading"):
+            ZipStore(MinimalReader(zip_bytes), mode=mode, read_only=False)  # type: ignore[arg-type]
+
+    def test_fsspec_file(self, tmp_path: Path, zip_bytes: bytes) -> None:
+        # a file opened through fsspec (already an io.IOBase) is used directly;
+        # fsspec's local filesystem stands in for a remote one
+        fsspec = pytest.importorskip("fsspec")
+
+        path = tmp_path / "fsspec.zip"
+        path.write_bytes(zip_bytes)
+        with fsspec.open(f"local://{path}", "rb") as fileobj:
+            store = ZipStore(fileobj, mode="r")
+            array = zarr.open_array(store, mode="r")
+            assert np.array_equal(array[...], np.arange(10))
+            assert store.path is None
+
+    def test_obstore_reader(self, tmp_path: Path, zip_bytes: bytes) -> None:
+        # obstore's ReadableFile is not an io.IOBase and its read() returns a
+        # buffer-protocol object; ZipStore adapts it via _RawReaderAdapter
+        obstore = pytest.importorskip("obstore")
+        from obstore.store import LocalStore as ObstoreLocalStore
+
+        (tmp_path / "obstore.zip").write_bytes(zip_bytes)
+        reader = obstore.open_reader(ObstoreLocalStore(str(tmp_path)), "obstore.zip")
+        store = ZipStore(reader, mode="r")
+        array = zarr.open_array(store, mode="r")
+        assert np.array_equal(array[...], np.arange(10))
+
+    def test_raw_reader_adapter_eof(self) -> None:
+        from zarr.storage._zip import _RawReaderAdapter
+
+        class MinimalReader:
+            """Non-io.IOBase reader exposing only read/seek/tell, like obstore."""
+
+            def __init__(self, data: bytes) -> None:
+                self._buffer = io.BytesIO(data)
+
+            def read(self, size: int, /) -> bytes:
+                return self._buffer.read(size)
+
+            def seek(self, pos: int, whence: int = 0, /) -> int:
+                return self._buffer.seek(pos, whence)
+
+            def tell(self) -> int:
+                return self._buffer.tell()
+
+        # the adapter must clamp reads to EOF: some readers (obstore < 0.6)
+        # raise on short reads instead of returning fewer bytes
+        data = b"0123456789"
+        adapter = _RawReaderAdapter(MinimalReader(data))  # type: ignore[arg-type]
+
+        # A read straddling EOF returns only the remaining bytes.
+        adapter.seek(len(data) - 3)
+        buf = bytearray(8)
+        assert adapter.readinto(buf) == 3
+        assert bytes(buf[:3]) == data[-3:]
+
+        # A read at EOF returns 0.
+        assert adapter.tell() == len(data)
+        assert adapter.readinto(bytearray(8)) == 0
+
+    def test_pickle_fileobj_raises(self, zip_bytes: bytes) -> None:
+        # an open file object cannot be reliably serialized, so pickling a
+        # file-object-backed store raises with a pointer at the alternative
+        store = ZipStore(io.BytesIO(zip_bytes), mode="r")
+        with pytest.raises(TypeError, match="cannot pickle a ZipStore backed by a file-like"):
+            pickle.dumps(store)
+
+    def test_pickle_path_backed_roundtrip(self, tmp_path: Path, zip_bytes: bytes) -> None:
+        # path-backed stores remain picklable: the path is serialized and the
+        # archive is reopened on unpickling
+        path = tmp_path / "pickled.zip"
+        path.write_bytes(zip_bytes)
+        store = ZipStore(path, mode="r")
+        unpickled = pickle.loads(pickle.dumps(store))
+        array = zarr.open_array(unpickled, mode="r")
+        assert np.array_equal(array[...], np.arange(10))
+
+    def test_str_and_eq(self, zip_bytes: bytes) -> None:
+        # file-object-backed stores stringify with the object repr and
+        # compare equal only when backed by the very same file object
+        fileobj = io.BytesIO(zip_bytes)
+        store = ZipStore(fileobj, mode="r")
+        assert str(store).startswith("zip://<")
+        assert store == ZipStore(fileobj, mode="r")
+        assert store != ZipStore(io.BytesIO(zip_bytes), mode="r")
+
+
 class ZipStoreLifecycleMachine(RuleBasedStateMachine):
     """Drive a ZipStore through construct / open / write / close transitions.