diff --git a/changes/4187.feature.md b/changes/4187.feature.md new file mode 100644 index 0000000000..87133e2034 --- /dev/null +++ b/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. diff --git a/docs/user-guide/storage.md b/docs/user-guide/storage.md index 7e0154b2a0..0ba6202c76 100644 --- a/docs/user-guide/storage.md +++ b/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 diff --git a/src/zarr/storage/_zip.py b/src/zarr/storage/_zip.py index 430b0c3e2a..69ae18bc2c 100644 --- a/src/zarr/storage/_zip.py +++ b/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,6 +212,10 @@ 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( @@ -137,13 +223,19 @@ async def clear(self) -> None: ) 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() diff --git a/tests/test_store/test_zip.py b/tests/test_store/test_zip.py index ed69114b51..0d8dadd18a 100644 --- a/tests/test_store/test_zip.py +++ b/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.