Update set_partial_values to accept Buffer rather than BytesLike

Author: maxrjones

src/zarr/abc/store.py

@@ -16,7 +16,6 @@
     from typing import Any, Self, TypeAlias
 
     from zarr.core.buffer import Buffer, BufferPrototype
-    from zarr.core.common import BytesLike
 
 __all__ = ["ByteGetter", "ByteSetter", "Store", "set_or_delete"]
 
@@ -285,9 +284,7 @@ def supports_partial_writes(self) -> bool:
         ...
 
     @abstractmethod
-    async def set_partial_values(
-        self, key_start_values: Iterable[tuple[str, int, BytesLike]]
-    ) -> None:
+    async def set_partial_values(self, key_start_values: Iterable[tuple[str, int, Buffer]]) -> None:
         """Store values at a given key, starting at byte range_start.
 
         Parameters

src/zarr/storage/_fsspec.py

@@ -18,7 +18,6 @@
     from fsspec.asyn import AsyncFileSystem
 
     from zarr.core.buffer import Buffer, BufferPrototype
-    from zarr.core.common import BytesLike
 
 
 ALLOWED_EXCEPTIONS: tuple[type[Exception], ...] = (
@@ -316,9 +315,7 @@ async def get_partial_values(
 
         return [None if isinstance(r, Exception) else prototype.buffer.from_bytes(r) for r in res]
 
-    async def set_partial_values(
-        self, key_start_values: Iterable[tuple[str, int, BytesLike]]
-    ) -> None:
+    async def set_partial_values(self, key_start_values: Iterable[tuple[str, int, Buffer]]) -> None:
         # docstring inherited
         raise NotImplementedError
 

src/zarr/storage/_local.py

@@ -174,7 +174,8 @@ async def _set(self, key: str, value: Buffer, exclusive: bool = False) -> None:
         await asyncio.to_thread(_put, path, value, start=None, exclusive=exclusive)
 
     async def set_partial_values(
-        self, key_start_values: Iterable[tuple[str, int, bytes | bytearray | memoryview]]
+        self,
+        key_start_values: Iterable[tuple[str, int, Buffer]],
     ) -> None:
         # docstring inherited
         self._check_writable()

src/zarr/storage/_logging.py

@@ -197,9 +197,7 @@ async def delete(self, key: str) -> None:
         with self.log(key):
             return await self._store.delete(key=key)
 
-    async def set_partial_values(
-        self, key_start_values: Iterable[tuple[str, int, bytes | bytearray | memoryview]]
-    ) -> None:
+    async def set_partial_values(self, key_start_values: Iterable[tuple[str, int, Buffer]]) -> None:
         # docstring inherited
         keys = ",".join([k[0] for k in key_start_values])
         with self.log(keys):

src/zarr/storage/_memory.py

@@ -38,7 +38,7 @@ class MemoryStore(Store):
 
     supports_writes: bool = True
     supports_deletes: bool = True
-    supports_partial_writes: bool = True
+    supports_partial_writes: bool = False
     supports_listing: bool = True
 
     _store_dict: MutableMapping[str, Buffer]
@@ -134,7 +134,7 @@ async def delete(self, key: str) -> None:
         except KeyError:
             logger.debug("Key %s does not exist.", key)
 
-    async def set_partial_values(self, key_start_values: Iterable[tuple[str, int, bytes]]) -> None:
+    async def set_partial_values(self, key_start_values: Iterable[tuple[str, int, Buffer]]) -> None:
         # docstring inherited
         raise NotImplementedError
 

src/zarr/storage/_wrapper.py

@@ -9,7 +9,6 @@
 
     from zarr.abc.store import ByteRequest
     from zarr.core.buffer import Buffer, BufferPrototype
-    from zarr.core.common import BytesLike
 
 from zarr.abc.store import Store
 
@@ -108,9 +107,7 @@ async def delete(self, key: str) -> None:
     def supports_partial_writes(self) -> bool:
         return self._store.supports_partial_writes
 
-    async def set_partial_values(
-        self, key_start_values: Iterable[tuple[str, int, BytesLike]]
-    ) -> None:
+    async def set_partial_values(self, key_start_values: Iterable[tuple[str, int, Buffer]]) -> None:
         return await self._store.set_partial_values(key_start_values)
 
     @property

src/zarr/storage/_zip.py

@@ -209,7 +209,7 @@ async def set(self, key: str, value: Buffer) -> None:
         with self._lock:
             self._set(key, value)
 
-    async def set_partial_values(self, key_start_values: Iterable[tuple[str, int, bytes]]) -> None:
+    async def set_partial_values(self, key_start_values: Iterable[tuple[str, int, Buffer]]) -> None:
         raise NotImplementedError
 
     async def set_if_not_exists(self, key: str, value: Buffer) -> None:

src/zarr/testing/store.py

@@ -181,6 +181,23 @@ async def test_set_many(self, store: S) -> None:
         for k, v in store_dict.items():
             assert (await self.get(store, k)).to_bytes() == v.to_bytes()
 
+    @pytest.mark.parametrize("key", ["zarr.json", "c/0", "foo/c/0.0", "foo/0/0"])
+    async def test_set_partial_values(self, store: S, key: str) -> None:
+        """
+        Ensure that data can be written to the store using the store.set method.
+        """
+        assert not store.read_only
+        # Create empty key
+        await store.set(key, self.buffer_cls.from_bytes(b""))
+        data_buf = self.buffer_cls.from_bytes(b"\x01\x02\x03\x04")
+        if store.supports_partial_writes:
+            await store.set_partial_values([(key, 0, data_buf[:2]), (key, 2, data_buf[2:])])
+            observed = await self.get(store, key)
+            assert_bytes_equal(observed.to_bytes(), data_buf)
+        else:
+            with pytest.raises(NotImplementedError):
+                await store.set_partial_values([(key, 0, data_buf[:2]), (key, 2, data_buf[2:])])
+
     @pytest.mark.parametrize(
         "key_ranges",
         [

tests/test_store/test_memory.py

@@ -41,7 +41,7 @@ def test_store_supports_listing(self, store: MemoryStore) -> None:
         assert store.supports_listing
 
     def test_store_supports_partial_writes(self, store: MemoryStore) -> None:
-        assert store.supports_partial_writes
+        assert not store.supports_partial_writes
 
     def test_list_prefix(self, store: MemoryStore) -> None:
         assert True