Backport PR #3138: Add with_read_only() convenience method to store (#3181)

Co-authored-by: Max Jones <[email protected]>

Author: meeseeksmachine

changes/3138.feature.rst

@@ -0,0 +1 @@
+Adds a `with_read_only` convenience method to the `Store` abstract base class (raises `NotImplementedError`) and implementations to the `MemoryStore`, `ObjectStore`, `LocalStore`, and `FsspecStore` classes.

src/zarr/abc/store.py

@@ -83,6 +83,27 @@ async def open(cls, *args: Any, **kwargs: Any) -> Self:
         await store._open()
         return store
 
+    def with_read_only(self, read_only: bool = False) -> Store:
+        """
+        Return a new store with a new read_only setting.
+
+        The new store points to the same location with the specified new read_only state.
+        The returned Store is not automatically opened, and this store is
+        not automatically closed.
+
+        Parameters
+        ----------
+        read_only
+            If True, the store will be created in read-only mode. Defaults to False.
+
+        Returns
+        -------
+            A new store of the same type with the new read only attribute.
+        """
+        raise NotImplementedError(
+            f"with_read_only is not implemented for the {type(self)} store type."
+        )
+
     def __enter__(self) -> Self:
         """Enter a context manager that will close the store upon exiting."""
         return self

src/zarr/storage/_fsspec.py

@@ -122,6 +122,7 @@ class FsspecStore(Store):
 
     fs: AsyncFileSystem
     allowed_exceptions: tuple[type[Exception], ...]
+    path: str
 
     def __init__(
         self,
@@ -258,6 +259,15 @@ def from_url(
 
         return cls(fs=fs, path=path, read_only=read_only, allowed_exceptions=allowed_exceptions)
 
+    def with_read_only(self, read_only: bool = False) -> FsspecStore:
+        # docstring inherited
+        return type(self)(
+            fs=self.fs,
+            path=self.path,
+            allowed_exceptions=self.allowed_exceptions,
+            read_only=read_only,
+        )
+
     async def clear(self) -> None:
         # docstring inherited
         try:

src/zarr/storage/_local.py

@@ -102,6 +102,13 @@ def __init__(self, root: Path | str, *, read_only: bool = False) -> None:
             )
         self.root = root
 
+    def with_read_only(self, read_only: bool = False) -> LocalStore:
+        # docstring inherited
+        return type(self)(
+            root=self.root,
+            read_only=read_only,
+        )
+
     async def _open(self) -> None:
         if not self.read_only:
             self.root.mkdir(parents=True, exist_ok=True)

src/zarr/storage/_memory.py

@@ -54,6 +54,13 @@ def __init__(
             store_dict = {}
         self._store_dict = store_dict
 
+    def with_read_only(self, read_only: bool = False) -> MemoryStore:
+        # docstring inherited
+        return type(self)(
+            store_dict=self._store_dict,
+            read_only=read_only,
+        )
+
     async def clear(self) -> None:
         # docstring inherited
         self._store_dict.clear()

src/zarr/storage/_obstore.py

@@ -69,6 +69,13 @@ def __init__(self, store: _UpstreamObjectStore, *, read_only: bool = False) -> N
         super().__init__(read_only=read_only)
         self.store = store
 
+    def with_read_only(self, read_only: bool = False) -> ObjectStore:
+        # docstring inherited
+        return type(self)(
+            store=self.store,
+            read_only=read_only,
+        )
+
     def __str__(self) -> str:
         return f"object_store://{self.store}"
 

src/zarr/testing/store.py

@@ -149,6 +149,58 @@ async def test_read_only_store_raises(self, open_kwargs: dict[str, Any]) -> None
         ):
             await store.delete("foo")
 
+    async def test_with_read_only_store(self, open_kwargs: dict[str, Any]) -> None:
+        kwargs = {**open_kwargs, "read_only": True}
+        store = await self.store_cls.open(**kwargs)
+        assert store.read_only
+
+        # Test that you cannot write to a read-only store
+        with pytest.raises(
+            ValueError, match="store was opened in read-only mode and does not support writing"
+        ):
+            await store.set("foo", self.buffer_cls.from_bytes(b"bar"))
+
+        # Check if the store implements with_read_only
+        try:
+            writer = store.with_read_only(read_only=False)
+        except NotImplementedError:
+            # Test that stores that do not implement with_read_only raise NotImplementedError with the correct message
+            with pytest.raises(
+                NotImplementedError,
+                match=f"with_read_only is not implemented for the {type(store)} store type.",
+            ):
+                store.with_read_only(read_only=False)
+            return
+
+        # Test that you can write to a new store copy
+        assert not writer._is_open
+        assert not writer.read_only
+        await writer.set("foo", self.buffer_cls.from_bytes(b"bar"))
+        await writer.delete("foo")
+
+        # Test that you cannot write to the original store
+        assert store.read_only
+        with pytest.raises(
+            ValueError, match="store was opened in read-only mode and does not support writing"
+        ):
+            await store.set("foo", self.buffer_cls.from_bytes(b"bar"))
+        with pytest.raises(
+            ValueError, match="store was opened in read-only mode and does not support writing"
+        ):
+            await store.delete("foo")
+
+        # Test that you cannot write to a read-only store copy
+        reader = store.with_read_only(read_only=True)
+        assert reader.read_only
+        with pytest.raises(
+            ValueError, match="store was opened in read-only mode and does not support writing"
+        ):
+            await reader.set("foo", self.buffer_cls.from_bytes(b"bar"))
+        with pytest.raises(
+            ValueError, match="store was opened in read-only mode and does not support writing"
+        ):
+            await reader.delete("foo")
+
     @pytest.mark.parametrize("key", ["c/0", "foo/c/0.0", "foo/0/0"])
     @pytest.mark.parametrize(
         ("data", "byte_range"),