From efeb4402f1cfe5bd25e4f435fb96b5b5516634d7 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 24 Sep 2024 13:13:19 -0500 Subject: [PATCH 01/14] Allow mode casting --- src/zarr/abc/store.py | 7 +++++++ src/zarr/store/common.py | 2 +- src/zarr/store/local.py | 5 +++++ src/zarr/store/memory.py | 5 +++++ src/zarr/store/remote.py | 9 +++++++++ src/zarr/store/zip.py | 7 +++++++ src/zarr/testing/store.py | 10 +++++++++- 7 files changed, 43 insertions(+), 2 deletions(-) diff --git a/src/zarr/abc/store.py b/src/zarr/abc/store.py index 2f02ac36a..32924e825 100644 --- a/src/zarr/abc/store.py +++ b/src/zarr/abc/store.py @@ -81,6 +81,13 @@ async def empty(self) -> bool: ... @abstractmethod async def clear(self) -> None: ... + @abstractmethod + def with_mode(self, mode: AccessModeLiteral) -> Self: + """ + Return a new store pointing to the same location with a new mode. + """ + ... + @property def mode(self) -> AccessMode: """Access mode of the store.""" diff --git a/src/zarr/store/common.py b/src/zarr/store/common.py index 0c126c63d..003c07f76 100644 --- a/src/zarr/store/common.py +++ b/src/zarr/store/common.py @@ -90,7 +90,7 @@ async def make_store_path( result = store_like elif isinstance(store_like, Store): if mode is not None: - assert AccessMode.from_literal(mode) == store_like.mode + store_like = store_like.with_mode(mode) await store_like._ensure_open() result = StorePath(store_like) elif store_like is None: diff --git a/src/zarr/store/local.py b/src/zarr/store/local.py index 39a94969e..11f487fbf 100644 --- a/src/zarr/store/local.py +++ b/src/zarr/store/local.py @@ -6,6 +6,8 @@ from pathlib import Path from typing import TYPE_CHECKING +from typing_extensions import Self + from zarr.abc.store import Store from zarr.core.buffer import Buffer from zarr.core.common import concurrent_map, to_thread @@ -103,6 +105,9 @@ async def empty(self) -> bool: else: return True + def with_mode(self, mode: AccessModeLiteral) -> Self: + return type(self)(root=self.root, mode=mode) + def __str__(self) -> str: return f"file://{self.root}" diff --git a/src/zarr/store/memory.py b/src/zarr/store/memory.py index 83734e894..ead88f4d2 100644 --- a/src/zarr/store/memory.py +++ b/src/zarr/store/memory.py @@ -3,6 +3,8 @@ from collections.abc import AsyncGenerator, MutableMapping from typing import TYPE_CHECKING +from typing_extensions import Self + from zarr.abc.store import Store from zarr.core.buffer import Buffer, gpu from zarr.core.common import concurrent_map @@ -42,6 +44,9 @@ async def empty(self) -> bool: async def clear(self) -> None: self._store_dict.clear() + def with_mode(self, mode: AccessModeLiteral) -> Self: + return type(self)(store_dict=self._store_dict, mode=mode) + def __str__(self) -> str: return f"memory://{id(self._store_dict)}" diff --git a/src/zarr/store/remote.py b/src/zarr/store/remote.py index 02bda6b1d..2622b0d6a 100644 --- a/src/zarr/store/remote.py +++ b/src/zarr/store/remote.py @@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Any import fsspec +from typing_extensions import Self from zarr.abc.store import Store from zarr.store.common import _dereference_path @@ -95,6 +96,14 @@ async def clear(self) -> None: async def empty(self) -> bool: return not await self.fs._find(self.path, withdirs=True) + def with_mode(self, mode: AccessModeLiteral) -> Self: + return type(self)( + fs=self.fs, + mode=mode, + path=self.path, + allowed_exceptions=self.allowed_exceptions, + ) + def __repr__(self) -> str: return f"" diff --git a/src/zarr/store/zip.py b/src/zarr/store/zip.py index 2e4927ace..4a2efe1d3 100644 --- a/src/zarr/store/zip.py +++ b/src/zarr/store/zip.py @@ -7,6 +7,8 @@ from pathlib import Path from typing import TYPE_CHECKING, Any, Literal +from typing_extensions import Self + from zarr.abc.store import Store from zarr.core.buffer import Buffer, BufferPrototype @@ -115,6 +117,11 @@ async def empty(self) -> bool: else: return True + def with_mode(self, mode: ZipStoreAccessModeLiteral) -> Self: # type: ignore[override] + return type(self)( + path=self.path, mode=mode, compression=self.compression, allowZip64=self.allowZip64 + ) + def __str__(self) -> str: return f"zip://{self.path}" diff --git a/src/zarr/testing/store.py b/src/zarr/testing/store.py index ebd4b85c9..4051ada61 100644 --- a/src/zarr/testing/store.py +++ b/src/zarr/testing/store.py @@ -1,10 +1,11 @@ import pickle -from typing import Any, Generic, TypeVar +from typing import Any, Generic, TypeVar, cast import pytest from zarr.abc.store import AccessMode, Store from zarr.core.buffer import Buffer, default_buffer_prototype +from zarr.core.common import AccessModeLiteral from zarr.core.sync import _collect_aiterator from zarr.store._utils import _normalize_interval_index from zarr.testing.utils import assert_bytes_equal @@ -251,3 +252,10 @@ async def test_list_dir(self, store: S) -> None: keys_observed = await _collect_aiterator(store.list_dir(root + "/")) assert sorted(keys_expected) == sorted(keys_observed) + + async def test_with_mode(self, store: S) -> None: + for mode in ["r", "w"]: + mode = cast(AccessModeLiteral, mode) + result = store.with_mode(mode) + assert result.mode == AccessMode.from_literal(mode) + assert isinstance(result, type(store)) From 501f151353cce5c656b3a56593ed3595d5740153 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Wed, 25 Sep 2024 13:17:55 -0500 Subject: [PATCH 02/14] fixup --- src/zarr/abc/store.py | 20 +++++++++++++++++++- src/zarr/store/zip.py | 4 +--- src/zarr/testing/store.py | 26 ++++++++++++++++++++++---- tests/v3/test_store/test_zip.py | 4 ++++ 4 files changed, 46 insertions(+), 8 deletions(-) diff --git a/src/zarr/abc/store.py b/src/zarr/abc/store.py index 32924e825..497bf7152 100644 --- a/src/zarr/abc/store.py +++ b/src/zarr/abc/store.py @@ -84,7 +84,25 @@ async def clear(self) -> None: ... @abstractmethod def with_mode(self, mode: AccessModeLiteral) -> Self: """ - Return a new store pointing to the same location with a new mode. + Return a new store of the same type pointing to the same location with a new mode. + + The returned Store is not automatically opened. Call :meth:`Store.open` before + using. + + Parameters + ---------- + mode: AccessModeLiteral + The new mode to use. + + Returns + ------- + store: + A new store of the same type with the new mode. + + Examples + -------- + >>> writer = zarr.store.MemoryStore(mode="w") + >>> reader = writer.with_mode("r") """ ... diff --git a/src/zarr/store/zip.py b/src/zarr/store/zip.py index 4a2efe1d3..456657ccc 100644 --- a/src/zarr/store/zip.py +++ b/src/zarr/store/zip.py @@ -118,9 +118,7 @@ async def empty(self) -> bool: return True def with_mode(self, mode: ZipStoreAccessModeLiteral) -> Self: # type: ignore[override] - return type(self)( - path=self.path, mode=mode, compression=self.compression, allowZip64=self.allowZip64 - ) + raise NotImplementedError("ZipStore cannot be reopened with a new mode.") def __str__(self) -> str: return f"zip://{self.path}" diff --git a/src/zarr/testing/store.py b/src/zarr/testing/store.py index 4051ada61..59af2e1df 100644 --- a/src/zarr/testing/store.py +++ b/src/zarr/testing/store.py @@ -254,8 +254,26 @@ async def test_list_dir(self, store: S) -> None: assert sorted(keys_expected) == sorted(keys_observed) async def test_with_mode(self, store: S) -> None: - for mode in ["r", "w"]: + data = b"0000" + self.set(store, "key", self.buffer_cls.from_bytes(data)) + assert self.get(store, "key").to_bytes() == data + + for mode in ["r", "a"]: mode = cast(AccessModeLiteral, mode) - result = store.with_mode(mode) - assert result.mode == AccessMode.from_literal(mode) - assert isinstance(result, type(store)) + clone = store.with_mode(mode) + # await store.close() + await clone._ensure_open() + assert clone.mode == AccessMode.from_literal(mode) + assert isinstance(clone, type(store)) + + # earlier writes are visible + assert self.get(clone, "key").to_bytes() == data + + # writes to original after with_mode is visible + self.set(store, "key-2", self.buffer_cls.from_bytes(data)) + assert self.get(clone, "key-2").to_bytes() == data + + if mode == "w": + # writes to clone is visible in the original + self.set(store, "key-3", self.buffer_cls.from_bytes(data)) + assert self.get(clone, "key-3").to_bytes() == data diff --git a/tests/v3/test_store/test_zip.py b/tests/v3/test_store/test_zip.py index 7c332e9a2..c95dcc77c 100644 --- a/tests/v3/test_store/test_zip.py +++ b/tests/v3/test_store/test_zip.py @@ -96,3 +96,7 @@ def test_api_integration(self, store: ZipStore) -> None: del root["bar"] store.close() + + async def test_with_mode(self, store: ZipStore) -> None: + with pytest.raises(NotImplementedError, match="new mode"): + await super().test_with_mode(store) From 2ef51f7b84cf26b08bed0c1faaba21933454eff2 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Wed, 25 Sep 2024 14:57:57 -0500 Subject: [PATCH 03/14] fixup --- src/zarr/store/memory.py | 18 ++++++++++++++---- tests/v3/test_store/test_memory.py | 5 +++++ 2 files changed, 19 insertions(+), 4 deletions(-) diff --git a/src/zarr/store/memory.py b/src/zarr/store/memory.py index ead88f4d2..1f28a61ff 100644 --- a/src/zarr/store/memory.py +++ b/src/zarr/store/memory.py @@ -17,6 +17,12 @@ from zarr.core.common import AccessModeLiteral +# T = TypeVar("T", bound=Buffer | gpu.Buffer) + + +# class _MemoryStore + + # TODO: this store could easily be extended to wrap any MutableMapping store from v2 # When that is done, the `MemoryStore` will just be a store that wraps a dict. class MemoryStore(Store): @@ -161,19 +167,23 @@ class GpuMemoryStore(MemoryStore): of the original location. This guarantees that chunks will always be in GPU memory for downstream processing. For location agnostic use cases, it would be better to use `MemoryStore` instead. + + Parameters + ---------- + store_dict: MutableMapping, optional + A mutable mapping with string keys and :class:`zarr.core.buffer.gpu.Buffer` + values. """ _store_dict: MutableMapping[str, Buffer] def __init__( self, - store_dict: MutableMapping[str, Buffer] | None = None, + store_dict: MutableMapping[str, gpu.Buffer] | None = None, *, mode: AccessModeLiteral = "r", ) -> None: - super().__init__(mode=mode) - if store_dict: - self._store_dict = {k: gpu.Buffer.from_buffer(store_dict[k]) for k in iter(store_dict)} + super().__init__(store_dict=store_dict, mode=mode) # type: ignore[arg-type] def __str__(self) -> str: return f"gpumemory://{id(self._store_dict)}" diff --git a/tests/v3/test_store/test_memory.py b/tests/v3/test_store/test_memory.py index 2498cdc24..1b9f56ae0 100644 --- a/tests/v3/test_store/test_memory.py +++ b/tests/v3/test_store/test_memory.py @@ -80,3 +80,8 @@ def test_store_supports_partial_writes(self, store: GpuMemoryStore) -> None: def test_list_prefix(self, store: GpuMemoryStore) -> None: assert True + + def test_dict_reference(self, store: GpuMemoryStore) -> None: + store_dict = {} + result = GpuMemoryStore(store_dict=store_dict) + assert result._store_dict is store_dict From 3ee37f1afc431d4eab2b2c11a27c3eeb91c03c2b Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Wed, 25 Sep 2024 15:12:19 -0500 Subject: [PATCH 04/14] fixup --- tests/v3/test_store/test_memory.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/v3/test_store/test_memory.py b/tests/v3/test_store/test_memory.py index 1b9f56ae0..4ae3371b7 100644 --- a/tests/v3/test_store/test_memory.py +++ b/tests/v3/test_store/test_memory.py @@ -58,9 +58,14 @@ def set(self, store: GpuMemoryStore, key: str, value: Buffer) -> None: def get(self, store: MemoryStore, key: str) -> Buffer: return store._store_dict[key] - @pytest.fixture(scope="function", params=[None, {}]) - def store_kwargs(self, request) -> dict[str, str | None | dict[str, Buffer]]: - return {"store_dict": request.param, "mode": "r+"} + @pytest.fixture(scope="function", params=[None, True]) + def store_kwargs( + self, request: pytest.FixtureRequest + ) -> dict[str, str | None | dict[str, Buffer]]: + kwargs = {"store_dict": None, "mode": "r+"} + if request.param is True: + kwargs["store_dict"] = {} + return kwargs @pytest.fixture(scope="function") def store(self, store_kwargs: str | None | dict[str, gpu.Buffer]) -> GpuMemoryStore: From 926c71acc9e8886581540088d6f457997181d8d6 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Thu, 26 Sep 2024 11:06:27 -0500 Subject: [PATCH 05/14] fixup --- src/zarr/store/memory.py | 39 ++++++++++++++++++++++-------- src/zarr/testing/store.py | 22 ++++++++++++----- tests/v3/test_store/test_memory.py | 9 +++++++ 3 files changed, 54 insertions(+), 16 deletions(-) diff --git a/src/zarr/store/memory.py b/src/zarr/store/memory.py index 49e6cdf01..9101f59ba 100644 --- a/src/zarr/store/memory.py +++ b/src/zarr/store/memory.py @@ -16,12 +16,6 @@ from zarr.core.common import AccessModeLiteral -# T = TypeVar("T", bound=Buffer | gpu.Buffer) - - -# class _MemoryStore - - # TODO: this store could easily be extended to wrap any MutableMapping store from v2 # When that is done, the `MemoryStore` will just be a store that wraps a dict. class MemoryStore(Store): @@ -163,9 +157,13 @@ async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]: class GpuMemoryStore(MemoryStore): """A GPU only memory store that stores every chunk in GPU memory irrespective - of the original location. This guarantees that chunks will always be in GPU - memory for downstream processing. For location agnostic use cases, it would - be better to use `MemoryStore` instead. + of the original location. + + The dictionary of buffers to initialize this memory store with *must* be + GPU Buffers. + + Writing data to this store through ``.set`` will move the buffer to the GPU + if necessary. Parameters ---------- @@ -174,7 +172,7 @@ class GpuMemoryStore(MemoryStore): values. """ - _store_dict: MutableMapping[str, Buffer] + _store_dict: MutableMapping[str, gpu.Buffer] # type: ignore[assignment] def __init__( self, @@ -190,6 +188,27 @@ def __str__(self) -> str: def __repr__(self) -> str: return f"GpuMemoryStore({str(self)!r})" + @classmethod + def from_dict(cls, store_dict: MutableMapping[str, Buffer]) -> Self: + """ + Create a GpuMemoryStore from a dictionary of buffers at any location. + + The dictionary backing the newly created ``GpuMemoryStore`` will not be + the same as ``store_dict``. + + Parameters + ---------- + store_dict: mapping + A mapping of strings keys to arbitrary Buffers. The buffer data + will be moved into a :class:`gpu.Buffer`. + + Returns + ------- + GpuMemoryStore + """ + gpu_store_dict = {k: gpu.Buffer.from_buffer(v) for k, v in store_dict.items()} + return cls(gpu_store_dict) + async def set(self, key: str, value: Buffer, byte_range: tuple[int, int] | None = None) -> None: self._check_writable() assert isinstance(key, str) diff --git a/src/zarr/testing/store.py b/src/zarr/testing/store.py index 8c324b3ab..677ed20d8 100644 --- a/src/zarr/testing/store.py +++ b/src/zarr/testing/store.py @@ -267,13 +267,23 @@ async def test_with_mode(self, store: S) -> None: assert isinstance(clone, type(store)) # earlier writes are visible - assert self.get(clone, "key").to_bytes() == data + result = await clone.get("key", default_buffer_prototype()) + assert result is not None + assert result.to_bytes() == data - # writes to original after with_mode is visible + # # writes to original after with_mode is visible self.set(store, "key-2", self.buffer_cls.from_bytes(data)) - assert self.get(clone, "key-2").to_bytes() == data + result = await clone.get("key-2", default_buffer_prototype()) + assert result is not None + assert result.to_bytes() == data - if mode == "w": + if mode == "a": # writes to clone is visible in the original - self.set(store, "key-3", self.buffer_cls.from_bytes(data)) - assert self.get(clone, "key-3").to_bytes() == data + await clone.set("key-3", self.buffer_cls.from_bytes(data)) + result = await clone.get("key-3", default_buffer_prototype()) + assert result is not None + assert result.to_bytes() == data + + else: + with pytest.raises(ValueError): + await clone.set("key-3", self.buffer_cls.from_bytes(data)) diff --git a/tests/v3/test_store/test_memory.py b/tests/v3/test_store/test_memory.py index 67850767d..efb61b332 100644 --- a/tests/v3/test_store/test_memory.py +++ b/tests/v3/test_store/test_memory.py @@ -90,3 +90,12 @@ def test_dict_reference(self, store: GpuMemoryStore) -> None: store_dict = {} result = GpuMemoryStore(store_dict=store_dict) assert result._store_dict is store_dict + + def test_from_dict(self): + d = { + "a": gpu.Buffer.from_bytes(b"aaaa"), + "b": cpu.Buffer.from_bytes(b"bbbb"), + } + result = GpuMemoryStore.from_dict(d) + for v in result._store_dict.values(): + assert type(v) is gpu.Buffer From 7f173d2e2a59f8e8615f1db5b88bb1495a46e03b Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Thu, 26 Sep 2024 11:07:22 -0500 Subject: [PATCH 06/14] match message --- src/zarr/testing/store.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/zarr/testing/store.py b/src/zarr/testing/store.py index 677ed20d8..62f69bec1 100644 --- a/src/zarr/testing/store.py +++ b/src/zarr/testing/store.py @@ -285,5 +285,5 @@ async def test_with_mode(self, store: S) -> None: assert result.to_bytes() == data else: - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="store mode"): await clone.set("key-3", self.buffer_cls.from_bytes(data)) From f2fc9c24d92a289737e6dfbc26e5f1f222b424a6 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Thu, 26 Sep 2024 11:10:36 -0500 Subject: [PATCH 07/14] Update src/zarr/testing/store.py Co-authored-by: Davis Bennett --- src/zarr/testing/store.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/zarr/testing/store.py b/src/zarr/testing/store.py index 62f69bec1..ce86f2fa4 100644 --- a/src/zarr/testing/store.py +++ b/src/zarr/testing/store.py @@ -271,7 +271,7 @@ async def test_with_mode(self, store: S) -> None: assert result is not None assert result.to_bytes() == data - # # writes to original after with_mode is visible + # writes to original after with_mode is visible self.set(store, "key-2", self.buffer_cls.from_bytes(data)) result = await clone.get("key-2", default_buffer_prototype()) assert result is not None From 92f6fec9cdd50b50c8fdeb1d3bbc18282c999569 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 27 Sep 2024 09:29:51 -0500 Subject: [PATCH 08/14] fixup --- tests/v3/test_store/test_logging.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/v3/test_store/test_logging.py b/tests/v3/test_store/test_logging.py index a263c2ae0..7bac6521a 100644 --- a/tests/v3/test_store/test_logging.py +++ b/tests/v3/test_store/test_logging.py @@ -5,6 +5,7 @@ import pytest import zarr +import zarr.store from zarr.core.buffer import default_buffer_prototype from zarr.store.logging import LoggingStore @@ -48,3 +49,10 @@ async def test_logging_store_counter(store: Store) -> None: assert wrapped.counter["list"] == 0 assert wrapped.counter["list_dir"] == 0 assert wrapped.counter["list_prefix"] == 0 + + +async def test_with_mode(): + wrapped = LoggingStore(store=zarr.store.MemoryStore(mode="w"), log_level="INFO") + new = await wrapped.with_mode(mode="r") + assert new.mode.str == "r" + assert new.log_level == "INFO" From 109afd8ced3b2179dddb36c1a114ac85a2bb5f66 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 27 Sep 2024 09:51:24 -0500 Subject: [PATCH 09/14] fixup --- src/zarr/store/logging.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/src/zarr/store/logging.py b/src/zarr/store/logging.py index 0c05b4265..6760068f0 100644 --- a/src/zarr/store/logging.py +++ b/src/zarr/store/logging.py @@ -5,7 +5,7 @@ import time from collections import defaultdict from contextlib import contextmanager -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Self from zarr.abc.store import AccessMode, ByteRangeRequest, Store @@ -13,6 +13,7 @@ from collections.abc import AsyncGenerator, Generator, Iterable from zarr.core.buffer import Buffer, BufferPrototype + from zarr.core.common import AccessModeLiteral class LoggingStore(Store): @@ -27,6 +28,8 @@ def __init__( ): self._store = store self.counter = defaultdict(int) + self.log_level = log_level + self.log_handler = log_handler self._configure_logger(log_level, log_handler) @@ -162,3 +165,10 @@ async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]: with self.log(): async for key in self._store.list_dir(prefix=prefix): yield key + + def with_mode(self, mode: AccessModeLiteral) -> Self: + return type(self)( + self._store.with_mode(mode), + log_level=self.log_level, + log_handler=self.log_handler, + ) From d1807203ca0a319eb6fe5292ffe70fc899b340f2 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 27 Sep 2024 12:48:15 -0500 Subject: [PATCH 10/14] fixup --- src/zarr/store/logging.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/zarr/store/logging.py b/src/zarr/store/logging.py index a74003a4b..36c5f6a8a 100644 --- a/src/zarr/store/logging.py +++ b/src/zarr/store/logging.py @@ -98,6 +98,12 @@ def _is_open(self) -> bool: # type: ignore[override] with self.log(): return self._store._is_open + async def _open(self) -> None: + return await self._store._open() + + async def _ensure_open(self) -> None: + return await self._store._ensure_open() + async def empty(self) -> bool: with self.log(): return await self._store.empty() From 3a02ee61982de445b45310e6042d56651ab4f04f Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 27 Sep 2024 12:53:05 -0500 Subject: [PATCH 11/14] fixup --- src/zarr/store/common.py | 2 +- tests/v3/test_store/test_logging.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/zarr/store/common.py b/src/zarr/store/common.py index a79564a50..08f94d3d1 100644 --- a/src/zarr/store/common.py +++ b/src/zarr/store/common.py @@ -89,7 +89,7 @@ async def make_store_path( assert AccessMode.from_literal(mode) == store_like.store.mode result = store_like elif isinstance(store_like, Store): - if mode is not None: + if mode is not None and mode != store_like.mode.str: store_like = store_like.with_mode(mode) await store_like._ensure_open() result = StorePath(store_like) diff --git a/tests/v3/test_store/test_logging.py b/tests/v3/test_store/test_logging.py index 7bac6521a..b03c9b94f 100644 --- a/tests/v3/test_store/test_logging.py +++ b/tests/v3/test_store/test_logging.py @@ -53,6 +53,6 @@ async def test_logging_store_counter(store: Store) -> None: async def test_with_mode(): wrapped = LoggingStore(store=zarr.store.MemoryStore(mode="w"), log_level="INFO") - new = await wrapped.with_mode(mode="r") + new = wrapped.with_mode(mode="r") assert new.mode.str == "r" assert new.log_level == "INFO" From a4f83fecd557be27cec74d79e6a2af08b1c8b11e Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 27 Sep 2024 15:21:09 -0500 Subject: [PATCH 12/14] pre-commit --- tests/v3/test_store/test_stateful_store.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/v3/test_store/test_stateful_store.py b/tests/v3/test_store/test_stateful_store.py index 1a43b55fe..efa1953a5 100644 --- a/tests/v3/test_store/test_stateful_store.py +++ b/tests/v3/test_store/test_stateful_store.py @@ -239,7 +239,7 @@ def check_zarr_keys(self) -> None: def test_zarr_hierarchy(sync_store: Store) -> None: - def mk_test_instance_sync(): + def mk_test_instance_sync() -> None: return ZarrStoreStateMachine(sync_store) if isinstance(sync_store, ZipStore): From a4dc888ebd5fcabbbe69a1781e95c5178c0785e3 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 27 Sep 2024 15:22:11 -0500 Subject: [PATCH 13/14] log methods --- src/zarr/store/logging.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/src/zarr/store/logging.py b/src/zarr/store/logging.py index 36c5f6a8a..d8dec25df 100644 --- a/src/zarr/store/logging.py +++ b/src/zarr/store/logging.py @@ -99,10 +99,12 @@ def _is_open(self) -> bool: # type: ignore[override] return self._store._is_open async def _open(self) -> None: - return await self._store._open() + with self.log(): + return await self._store._open() async def _ensure_open(self) -> None: - return await self._store._ensure_open() + with self.log(): + return await self._store._ensure_open() async def empty(self) -> bool: with self.log(): @@ -173,8 +175,9 @@ async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]: yield key def with_mode(self, mode: AccessModeLiteral) -> Self: - return type(self)( - self._store.with_mode(mode), - log_level=self.log_level, - log_handler=self.log_handler, - ) + with self.log(): + return type(self)( + self._store.with_mode(mode), + log_level=self.log_level, + log_handler=self.log_handler, + ) From b302a0fde8bfb461c50dae62daf4d4b3e37dea69 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 27 Sep 2024 21:02:18 +0000 Subject: [PATCH 14/14] style: pre-commit fixes --- src/zarr/testing/store.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/zarr/testing/store.py b/src/zarr/testing/store.py index a739bdb59..5495e6fdf 100644 --- a/src/zarr/testing/store.py +++ b/src/zarr/testing/store.py @@ -275,7 +275,6 @@ async def test_list_dir(self, store: S) -> None: keys_observed = await _collect_aiterator(store.list_dir(root + "/")) assert sorted(keys_expected) == sorted(keys_observed) - async def test_with_mode(self, store: S) -> None: data = b"0000" self.set(store, "key", self.buffer_cls.from_bytes(data)) @@ -326,4 +325,3 @@ async def test_set_if_not_exists(self, store: S) -> None: result = await store.get("k2", default_buffer_prototype()) assert result == new -