diff --git a/changes/4450.bugfix.md b/changes/4450.bugfix.md new file mode 100644 index 0000000000..ea68727f33 --- /dev/null +++ b/changes/4450.bugfix.md @@ -0,0 +1 @@ +`ZipStore` now opens its archive on first use from every method, so calling `get`, `get_partial_values`, `set_if_not_exists`, or `clear` on a store that was never opened no longer raises `AttributeError`. A store created with `mode="w"` or `mode="x"` now keeps its entries when it is reopened by `move()` or by unpickling, instead of truncating the archive or refusing to open it. diff --git a/src/zarr/storage/_zip.py b/src/zarr/storage/_zip.py index 6b9ebca561..9e773ae9fe 100644 --- a/src/zarr/storage/_zip.py +++ b/src/zarr/storage/_zip.py @@ -162,22 +162,33 @@ def __init__( self._zmode = mode self.compression = compression self.allowZip64 = allowZip64 + self._lock = threading.RLock() def _sync_open(self) -> None: if self._is_open: raise ValueError("store is already open") - self._lock = threading.RLock() - self._zf = zipfile.ZipFile( self.path if self.path is not None else self._fileobj, # type: ignore[arg-type] mode=self._zmode, compression=self.compression, allowZip64=self.allowZip64, ) + # "w" truncates and "x" refuses an existing file. Both apply only to the + # first open: reopening after close(), move(), or unpickling must keep + # the entries already written. + if self._zmode in ("w", "x"): + self._zmode = "a" self._is_open = True + def _zipfile(self) -> zipfile.ZipFile: + """Return the archive, opening it on first use.""" + with self._lock: + if not self._is_open: + self._sync_open() + return self._zf + async def _open(self) -> None: self._sync_open() @@ -197,6 +208,7 @@ def __getstate__(self) -> dict[str, Any]: def __setstate__(self, state: dict[str, Any]) -> None: self.__dict__ = state + self._lock = threading.RLock() self._is_open = False self._sync_open() @@ -216,7 +228,7 @@ async def clear(self) -> None: raise NotImplementedError( "clear() is not supported for a ZipStore backed by a file-like object" ) - self._zf.close() + self._zipfile().close() os.remove(self.path) self._zf = zipfile.ZipFile( self.path, mode="w", compression=self.compression, allowZip64=self.allowZip64 @@ -243,11 +255,9 @@ def _get( prototype: BufferPrototype, byte_range: ByteRequest | None = None, ) -> Buffer | None: - if not self._is_open: - self._sync_open() # docstring inherited try: - with self._zf.open(key) as f: # will raise KeyError + with self._zipfile().open(key) as f: # will raise KeyError if byte_range is None: return prototype.buffer.from_bytes(f.read()) elif isinstance(byte_range, RangeByteRequest): @@ -288,8 +298,6 @@ async def get_partial_values( return out def _set(self, key: str, value: Buffer) -> None: - if not self._is_open: - self._sync_open() # generally, this should be called inside a lock keyinfo = zipfile.ZipInfo(filename=key, date_time=time.localtime(time.time())[:6]) keyinfo.compress_type = self.compression @@ -298,13 +306,11 @@ def _set(self, key: str, value: Buffer) -> None: keyinfo.external_attr |= 0x10 # MS-DOS directory flag else: keyinfo.external_attr = 0o644 << 16 # ?rw-r--r-- - self._zf.writestr(keyinfo, value.to_bytes()) + self._zipfile().writestr(keyinfo, value.to_bytes()) async def set(self, key: str, value: Buffer) -> None: # docstring inherited self._check_writable() - if not self._is_open: - self._sync_open() if not isinstance(value, Buffer): raise TypeError( f"ZipStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead." @@ -315,7 +321,7 @@ async def set(self, key: str, value: Buffer) -> None: async def set_if_not_exists(self, key: str, value: Buffer) -> None: self._check_writable() with self._lock: - members = self._zf.namelist() + members = self._zipfile().namelist() if key not in members: self._set(key, value) @@ -337,11 +343,9 @@ async def delete(self, key: str) -> None: async def exists(self, key: str) -> bool: # docstring inherited - if not self._is_open: - self._sync_open() with self._lock: try: - self._zf.getinfo(key) + self._zipfile().getinfo(key) except KeyError: return False else: @@ -349,10 +353,8 @@ async def exists(self, key: str) -> bool: async def list(self) -> AsyncIterator[str]: # docstring inherited - if not self._is_open: - self._sync_open() with self._lock: - for key in self._zf.namelist(): + for key in self._zipfile().namelist(): yield key async def list_prefix(self, prefix: str) -> AsyncIterator[str]: @@ -363,11 +365,9 @@ async def list_prefix(self, prefix: str) -> AsyncIterator[str]: async def list_dir(self, prefix: str) -> AsyncIterator[str]: # docstring inherited - if not self._is_open: - self._sync_open() prefix = prefix.rstrip("/") - keys = self._zf.namelist() + keys = self._zipfile().namelist() seen = set() if prefix == "": keys_unique = {k.split("/")[0] for k in keys} diff --git a/tests/test_store/test_zip.py b/tests/test_store/test_zip.py index 32b18c5273..1189fbe6d2 100644 --- a/tests/test_store/test_zip.py +++ b/tests/test_store/test_zip.py @@ -188,6 +188,76 @@ async def test_move(self, tmp_path: Path) -> None: assert not origin.exists() assert np.array_equal(array[...], np.arange(10)) + @pytest.mark.parametrize( + "call", + [ + lambda s: s.get("foo", default_buffer_prototype()), + lambda s: s.get_partial_values(default_buffer_prototype(), [("foo", None)]), + lambda s: s.exists("foo"), + lambda s: s.set("bar", cpu.Buffer.from_bytes(b"x")), + lambda s: s.set_if_not_exists("bar", cpu.Buffer.from_bytes(b"x")), + lambda s: s.clear(), + lambda s: s.delete("bar"), + lambda s: s.delete_dir("bar"), + lambda s: s.is_empty(""), + ], + ids=[ + "get", + "get_partial_values", + "exists", + "set", + "set_if_not_exists", + "clear", + "delete", + "delete_dir", + "is_empty", + ], + ) + async def test_methods_open_store_on_first_use( + self, store_kwargs: dict[str, Any], call: Any + ) -> None: + # every method works on a store that was constructed but never opened + seed = await self.store_cls.open(**store_kwargs) + await seed.set("foo", cpu.Buffer.from_bytes(b"bar")) + seed.close() + + store = self.store_cls(**{**store_kwargs, "mode": "a"}) + assert not store._is_open + await call(store) + store.close() + + @pytest.mark.parametrize("method", ["list", "list_dir"]) + async def test_listing_opens_store_on_first_use( + self, store_kwargs: dict[str, Any], method: str + ) -> None: + seed = await self.store_cls.open(**store_kwargs) + await seed.set("foo", cpu.Buffer.from_bytes(b"bar")) + seed.close() + + store = self.store_cls(**{**store_kwargs, "mode": "r"}) + args = () if method == "list" else ("",) + assert [k async for k in getattr(store, method)(*args)] == ["foo"] + store.close() + + @pytest.mark.parametrize("mode", ["w", "x"]) + @pytest.mark.parametrize("reopen", ["close_twice", "move", "pickle"]) + async def test_reopen_keeps_entries(self, tmp_path: Path, mode: str, reopen: str) -> None: + # "w" truncates and "x" refuses an existing file; neither may apply + # when a store that already wrote entries is opened again + store = ZipStore(tmp_path / "data.zip", mode=mode) # type: ignore[arg-type] + await store.set("foo", cpu.Buffer.from_bytes(b"bar")) + if reopen == "close_twice": + store.close() + store.close() + elif reopen == "move": + await store.move(tmp_path / "moved" / "data.zip") + store.close() + else: + store.close() + pickle.loads(pickle.dumps(store)).close() + with zipfile.ZipFile(store.path) as zf: # type: ignore[arg-type] + assert zf.namelist() == ["foo"] + class TestZipStoreFileObj: """ZipStore backed by an open binary file-like object instead of a path.""" @@ -352,8 +422,7 @@ class ZipStoreLifecycleMachine(RuleBasedStateMachine): Invariant under test: a constructed ZipStore can always be closed without raising, regardless of whether it was ever opened or did any I/O. This is a property-based generalization of the former example-based regression tests - for ZipStore.close() being called on a never-opened store (which raised - AttributeError because ``_lock`` is created lazily in ``_sync_open``). + for ZipStore.close() being called on a never-opened store. """ def __init__(self, tmp_path: Path) -> None: