Allow mode casting

Author: TomAugspurger

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."""

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:

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}"
 

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)}"
 

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"<RemoteStore({type(self.fs).__name__}, {self.path})>"
 

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}"
 

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))