Support with_read_only in LoggingStore and LatencyStore.

Fixes #3699

Author: BorisTheBrave

src/zarr/experimental/cache_store.py

@@ -4,7 +4,7 @@
 import logging
 import time
 from collections import OrderedDict
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any, Literal, Self
 
 from zarr.abc.store import ByteRequest, Store
 from zarr.storage._wrapper import WrapperStore
@@ -120,6 +120,11 @@ def __init__(
         self._misses = 0
         self._evictions = 0
 
+    def _with_store(self, store: Store) -> Self:
+        # Cannot support this operation because it would share a cache, but have a new store
+        # So cache keys would conflict
+        raise NotImplementedError("CacheStore does not support this operation.")
+
     def _is_key_fresh(self, key: str) -> bool:
         """Check if a cached key is still fresh based on max_age_seconds.
 

src/zarr/storage/_logging.py

@@ -77,6 +77,9 @@ def _default_handler(self) -> logging.Handler:
         )
         return handler
 
+    def _with_store(self, store: T_Store) -> Self:
+        return type(self)(store=store, log_level=self.log_level, log_handler=self.log_handler)
+
     @contextmanager
     def log(self, hint: Any = "") -> Generator[None, None, None]:
         """Context manager to log method calls

src/zarr/storage/_wrapper.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Generic, TypeVar
+from typing import TYPE_CHECKING, Generic, TypeVar, cast
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, AsyncIterator, Iterable
@@ -31,14 +31,23 @@ class WrapperStore(Store, Generic[T_Store]):
     def __init__(self, store: T_Store) -> None:
         self._store = store
 
+    def _with_store(self, store: T_Store) -> Self:
+        """
+        Constructs a new instance of the wrapper store with the same details but a new store.
+        """
+        return type(self)(store=store)
+
     @classmethod
     async def open(cls: type[Self], store_cls: type[T_Store], *args: Any, **kwargs: Any) -> Self:
         store = store_cls(*args, **kwargs)
         await store._open()
         return cls(store=store)
 
+    def with_read_only(self, read_only: bool = False) -> Self:
+        return self._with_store(cast(T_Store, self._store.with_read_only(read_only)))
+
     def __enter__(self) -> Self:
-        return type(self)(self._store.__enter__())
+        return self._with_store(self._store.__enter__())
 
     def __exit__(
         self,

src/zarr/testing/store.py

@@ -4,7 +4,7 @@
 import json
 import pickle
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Generic, TypeVar
+from typing import TYPE_CHECKING, Generic, Self, TypeVar
 
 from zarr.storage import WrapperStore
 
@@ -578,10 +578,13 @@ class LatencyStore(WrapperStore[Store]):
     get_latency: float
     set_latency: float
 
-    def __init__(self, cls: Store, *, get_latency: float = 0, set_latency: float = 0) -> None:
+    def __init__(self, store: Store, *, get_latency: float = 0, set_latency: float = 0) -> None:
         self.get_latency = float(get_latency)
         self.set_latency = float(set_latency)
-        self._store = cls
+        self._store = store
+
+    def _with_store(self, store: Store) -> Self:
+        return type(self)(store, get_latency=self.get_latency, set_latency=self.set_latency)
 
     async def set(self, key: str, value: Buffer) -> None:
         """

tests/test_store/test_latency.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+import pytest
+
+from zarr.core.buffer import default_buffer_prototype
+from zarr.storage import MemoryStore
+from zarr.testing.store import LatencyStore
+
+
+async def test_latency_store_with_read_only_round_trip() -> None:
+    """
+    Ensure that LatencyStore.with_read_only returns another LatencyStore with
+    the requested read_only state, preserves latency configuration, and does
+    not change the original wrapper.
+    """
+    base = await MemoryStore.open()
+    # Start from a read-only underlying store
+    ro_base = base.with_read_only(read_only=True)
+    latency_ro = LatencyStore(ro_base, get_latency=0.01, set_latency=0.02)
+
+    assert latency_ro.read_only
+    assert latency_ro.get_latency == pytest.approx(0.01)
+    assert latency_ro.set_latency == pytest.approx(0.02)
+
+    buf = default_buffer_prototype().buffer.from_bytes(b"abcd")
+
+    # Cannot write through the read-only wrapper
+    with pytest.raises(
+        ValueError, match="store was opened in read-only mode and does not support writing"
+    ):
+        await latency_ro.set("key", buf)
+
+    # Create a writable wrapper from the read-only one
+    writer = latency_ro.with_read_only(read_only=False)
+    assert isinstance(writer, LatencyStore)
+    assert not writer.read_only
+    # Latency configuration is preserved
+    assert writer.get_latency == latency_ro.get_latency
+    assert writer.set_latency == latency_ro.set_latency
+
+    # Writes via the writable wrapper succeed
+    await writer.set("key", buf)
+    out = await writer.get("key", prototype=default_buffer_prototype())
+    assert out is not None
+    assert out.to_bytes() == buf.to_bytes()
+
+    # Creating a read-only copy from the writable wrapper works and is enforced
+    reader = writer.with_read_only(read_only=True)
+    assert isinstance(reader, LatencyStore)
+    assert reader.read_only
+    with pytest.raises(
+        ValueError, match="store was opened in read-only mode and does not support writing"
+    ):
+        await reader.set("other", buf)
+
+    # The original read-only wrapper remains read-only
+    assert latency_ro.read_only

tests/test_store/test_logging.py

@@ -86,6 +86,46 @@ def test_is_open_setter_raises(self, store: LoggingStore[LocalStore]) -> None:
         ):
             store._is_open = True
 
+    async def test_with_read_only_round_trip(self, local_store: LocalStore) -> None:
+        """
+        Ensure that LoggingStore.with_read_only returns another LoggingStore with
+        the requested read_only state, preserves logging configuration, and does
+        not change the original store.
+        """
+        # Start from a read-only underlying store
+        ro_store = local_store.with_read_only(read_only=True)
+        wrapped_ro = LoggingStore(store=ro_store, log_level="INFO")
+        assert wrapped_ro.read_only
+
+        buf = default_buffer_prototype().buffer.from_bytes(b"0123")
+
+        # Cannot write through the read-only wrapper
+        with pytest.raises(
+            ValueError, match="store was opened in read-only mode and does not support writing"
+        ):
+            await wrapped_ro.set("foo", buf)
+
+        # Create a writable wrapper
+        writer = wrapped_ro.with_read_only(read_only=False)
+        assert isinstance(writer, LoggingStore)
+        assert not writer.read_only
+        # logging configuration is preserved
+        assert writer.log_level == wrapped_ro.log_level
+        assert writer.log_handler == wrapped_ro.log_handler
+
+        # Writes via the writable wrapper succeed
+        await writer.set("foo", buf)
+        out = await writer.get("foo", prototype=default_buffer_prototype())
+        assert out is not None
+        assert out.to_bytes() == buf.to_bytes()
+
+        # The original wrapper remains read-only
+        assert wrapped_ro.read_only
+        with pytest.raises(
+            ValueError, match="store was opened in read-only mode and does not support writing"
+        ):
+            await wrapped_ro.set("bar", buf)
+
 
 @pytest.mark.parametrize("store", ["local", "memory", "zip"], indirect=["store"])
 async def test_logging_store(store: Store, caplog: pytest.LogCaptureFixture) -> None: