add BufferLike as buffer parameter for store methods that allocate memory

Author: d-v-b

src/zarr/abc/store.py

@@ -7,16 +7,17 @@
 from itertools import starmap
 from typing import TYPE_CHECKING, Literal, Protocol, runtime_checkable
 
+from zarr.core.buffer import Buffer, BufferPrototype
 from zarr.core.sync import sync
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, AsyncIterator, Iterable
     from types import TracebackType
     from typing import Any, Self, TypeAlias
 
-    from zarr.core.buffer import Buffer, BufferPrototype
+__all__ = ["BufferLike", "ByteGetter", "ByteSetter", "Store", "set_or_delete"]
 
-__all__ = ["ByteGetter", "ByteSetter", "Store", "set_or_delete"]
+BufferLike = type[Buffer] | BufferPrototype
 
 
 @dataclass
@@ -183,20 +184,31 @@ def __eq__(self, value: object) -> bool:
         """Equality comparison."""
         ...
 
+    @abstractmethod
+    def _get_default_buffer_class(self) -> type[Buffer]:
+        """
+        Get the default buffer class for this store.
+        """
+        ...
+
     @abstractmethod
     async def get(
         self,
         key: str,
-        prototype: BufferPrototype,
+        prototype: BufferLike | None = None,
         byte_range: ByteRequest | None = None,
     ) -> Buffer | None:
         """Retrieve the value associated with a given key.
 
         Parameters
         ----------
         key : str
-        prototype : BufferPrototype
-            The prototype of the output buffer. Stores may support a default buffer prototype.
+        prototype : BufferLike | None, optional
+            The prototype of the output buffer.
+            Can be either a Buffer class or an instance of `BufferPrototype`, in which the
+            `buffer` attribute will be used.
+            If `None`, the default buffer class for this store will be retrieved via the
+            ``_get_default_buffer_class`` method.
         byte_range : ByteRequest, optional
             ByteRequest may be one of the following. If not provided, all data associated with the key is retrieved.
             - RangeByteRequest(int, int): Request a specific range of bytes in the form (start, end). The end is exclusive. If the given range is zero-length or starts after the end of the object, an error will be returned. Additionally, if the range ends after the end of the object, the entire remainder of the object will be returned. Otherwise, the exact requested range will be returned.
@@ -210,7 +222,11 @@ async def get(
         ...
 
     async def get_bytes(
-        self, key: str, *, prototype: BufferPrototype, byte_range: ByteRequest | None = None
+        self,
+        key: str,
+        *,
+        prototype: BufferLike | None = None,
+        byte_range: ByteRequest | None = None,
     ) -> bytes:
         """
         Retrieve raw bytes from the store asynchronously.
@@ -222,8 +238,12 @@ async def get_bytes(
         ----------
         key : str
             The key identifying the data to retrieve.
-        prototype : BufferPrototype
-            The buffer prototype to use for reading the data.
+        prototype : BufferLike | None, optional
+            The prototype of the output buffer.
+            Can be either a Buffer class or an instance of `BufferPrototype`, in which the
+            `buffer` attribute will be used.
+            If `None`, the default buffer prototype for this store will be retrieved via the
+            ``_get_default_buffer_class`` method.
         byte_range : ByteRequest, optional
             If specified, only retrieve a portion of the stored data.
             Can be a ``RangeByteRequest``, ``OffsetByteRequest``, or ``SuffixByteRequest``.
@@ -258,7 +278,11 @@ async def get_bytes(
         return buffer.to_bytes()
 
     def get_bytes_sync(
-        self, key: str = "", *, prototype: BufferPrototype, byte_range: ByteRequest | None = None
+        self,
+        key: str = "",
+        *,
+        prototype: BufferLike | None = None,
+        byte_range: ByteRequest | None = None,
     ) -> bytes:
         """
         Retrieve raw bytes from the store synchronously.
@@ -271,8 +295,12 @@ def get_bytes_sync(
         ----------
         key : str, optional
             The key identifying the data to retrieve. Defaults to an empty string.
-        prototype : BufferPrototype
-            The buffer prototype to use for reading the data.
+        prototype : BufferLike | None, optional
+            The prototype of the output buffer.
+            Can be either a Buffer class or an instance of `BufferPrototype`, in which the
+            `buffer` attribute will be used.
+            If `None`, the default buffer prototype for this store will be retrieved via the
+            ``_get_default_buffer_class`` method.
         byte_range : ByteRequest, optional
             If specified, only retrieve a portion of the stored data.
             Can be a ``RangeByteRequest``, ``OffsetByteRequest``, or ``SuffixByteRequest``.
@@ -309,7 +337,11 @@ def get_bytes_sync(
         return sync(self.get_bytes(key, prototype=prototype, byte_range=byte_range))
 
     async def get_json(
-        self, key: str, *, prototype: BufferPrototype, byte_range: ByteRequest | None = None
+        self,
+        key: str,
+        *,
+        prototype: BufferLike | None = None,
+        byte_range: ByteRequest | None = None,
     ) -> Any:
         """
         Retrieve and parse JSON data from the store asynchronously.
@@ -321,8 +353,12 @@ async def get_json(
         ----------
         key : str
             The key identifying the JSON data to retrieve.
-        prototype : BufferPrototype
-            The buffer prototype to use for reading the data.
+        prototype : BufferLike | None, optional
+            The prototype of the output buffer.
+            Can be either a Buffer class or an instance of `BufferPrototype`, in which the
+            `buffer` attribute will be used.
+            If `None`, the default buffer prototype for this store will be retrieved via the
+            ``_get_default_buffer_class`` method.
         byte_range : ByteRequest, optional
             If specified, only retrieve a portion of the stored data.
             Can be a ``RangeByteRequest``, ``OffsetByteRequest``, or ``SuffixByteRequest``.
@@ -359,7 +395,11 @@ async def get_json(
         return json.loads(await self.get_bytes(key, prototype=prototype, byte_range=byte_range))
 
     def get_json_sync(
-        self, key: str = "", *, prototype: BufferPrototype, byte_range: ByteRequest | None = None
+        self,
+        key: str = "",
+        *,
+        prototype: BufferLike | None = None,
+        byte_range: ByteRequest | None = None,
     ) -> Any:
         """
         Retrieve and parse JSON data from the store synchronously.
@@ -372,8 +412,12 @@ def get_json_sync(
         ----------
         key : str, optional
             The key identifying the JSON data to retrieve. Defaults to an empty string.
-        prototype : BufferPrototype
-            The buffer prototype to use for reading the data.
+        prototype : BufferLike | None, optional
+            The prototype of the output buffer.
+            Can be either a Buffer class or an instance of `BufferPrototype`, in which the
+            `buffer` attribute will be used.
+            If `None`, the default buffer prototype for this store will be retrieved via the
+            ``_get_default_buffer_class`` method.
         byte_range : ByteRequest, optional
             If specified, only retrieve a portion of the stored data.
             Can be a ``RangeByteRequest``, ``OffsetByteRequest``, or ``SuffixByteRequest``.
@@ -417,15 +461,19 @@ def get_json_sync(
     @abstractmethod
     async def get_partial_values(
         self,
-        prototype: BufferPrototype,
+        prototype: BufferLike | None,
         key_ranges: Iterable[tuple[str, ByteRequest | None]],
     ) -> list[Buffer | None]:
         """Retrieve possibly partial values from given key_ranges.
 
         Parameters
         ----------
-        prototype : BufferPrototype
-            The prototype of the output buffer. Stores may support a default buffer prototype.
+        prototype : BufferLike | None
+            The prototype of the output buffer.
+            Can be either a Buffer class or an instance of `BufferPrototype`, in which the
+            `buffer` attribute will be used.
+            If `None`, the default buffer prototype for this store will be retrieved via the
+            ``_get_default_buffer_class`` method.
         key_ranges : Iterable[tuple[str, tuple[int | None, int | None]]]
             Ordered set of key, range pairs, a key may occur multiple times with different ranges
 
@@ -597,7 +645,7 @@ def close(self) -> None:
         self._is_open = False
 
     async def _get_many(
-        self, requests: Iterable[tuple[str, BufferPrototype, ByteRequest | None]]
+        self, requests: Iterable[tuple[str, BufferLike | None, ByteRequest | None]]
     ) -> AsyncGenerator[tuple[str, Buffer | None], None]:
         """
         Retrieve a collection of objects from storage. In general this method does not guarantee

src/zarr/experimental/cache_store.py

@@ -6,13 +6,13 @@
 from collections import OrderedDict
 from typing import TYPE_CHECKING, Any, Literal
 
-from zarr.abc.store import ByteRequest, Store
+from zarr.abc.store import BufferLike, ByteRequest, Store
 from zarr.storage._wrapper import WrapperStore
 
 logger = logging.getLogger(__name__)
 
 if TYPE_CHECKING:
-    from zarr.core.buffer.core import Buffer, BufferPrototype
+    from zarr.core.buffer.core import Buffer
 
 
 class CacheStore(WrapperStore[Store]):
@@ -218,7 +218,7 @@ def _remove_from_tracking(self, key: str) -> None:
         self._key_sizes.pop(key, None)
 
     async def _get_try_cache(
-        self, key: str, prototype: BufferPrototype, byte_range: ByteRequest | None = None
+        self, key: str, prototype: BufferLike | None, byte_range: ByteRequest | None = None
     ) -> Buffer | None:
         """Try to get data from cache first, falling back to source store."""
         maybe_cached_result = await self._cache.get(key, prototype, byte_range)
@@ -246,7 +246,7 @@ async def _get_try_cache(
             return maybe_fresh_result
 
     async def _get_no_cache(
-        self, key: str, prototype: BufferPrototype, byte_range: ByteRequest | None = None
+        self, key: str, prototype: BufferLike | None, byte_range: ByteRequest | None = None
     ) -> Buffer | None:
         """Get data directly from source store and update cache."""
         self._misses += 1
@@ -265,7 +265,7 @@ async def _get_no_cache(
     async def get(
         self,
         key: str,
-        prototype: BufferPrototype,
+        prototype: BufferLike | None = None,
         byte_range: ByteRequest | None = None,
     ) -> Buffer | None:
         """
@@ -275,8 +275,12 @@ async def get(
         ----------
         key : str
             The key to retrieve
-        prototype : BufferPrototype
-            Buffer prototype for creating the result buffer
+        prototype : BufferLike | None, optional
+            The prototype of the output buffer.
+            Can be either a Buffer class or an instance of `BufferPrototype`, in which the
+            `buffer` attribute will be used.
+            If `None`, the default buffer class for this store will be retrieved via the
+            ``_get_default_buffer_class`` method.
         byte_range : ByteRequest, optional
             Byte range to retrieve
 

src/zarr/storage/_common.py

@@ -3,10 +3,10 @@
 import importlib.util
 import json
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Literal, Self, TypeAlias
+from typing import Any, Literal, Self, TypeAlias
 
-from zarr.abc.store import ByteRequest, Store
-from zarr.core.buffer import Buffer, default_buffer_prototype
+from zarr.abc.store import BufferLike, ByteRequest, Store
+from zarr.core.buffer import Buffer
 from zarr.core.common import (
     ANY_ACCESS_MODE,
     ZARR_JSON,
@@ -26,9 +26,6 @@
 else:
     FSMap = None
 
-if TYPE_CHECKING:
-    from zarr.core.buffer import BufferPrototype
-
 
 def _dereference_path(root: str, path: str) -> str:
     if not isinstance(root, str):
@@ -145,16 +142,20 @@ async def open(cls, store: Store, path: str, mode: AccessModeLiteral | None = No
 
     async def get(
         self,
-        prototype: BufferPrototype | None = None,
+        prototype: BufferLike | None = None,
         byte_range: ByteRequest | None = None,
     ) -> Buffer | None:
         """
         Read bytes from the store.
 
         Parameters
         ----------
-        prototype : BufferPrototype, optional
-            The buffer prototype to use when reading the bytes.
+        prototype : BufferLike | None, optional
+            The prototype of the output buffer.
+            Can be either a Buffer class or an instance of `BufferPrototype`, in which the
+            `buffer` attribute will be used.
+            If `None`, the default buffer class for this store will be retrieved via the
+            store's ``_get_default_buffer_class`` method.
         byte_range : ByteRequest, optional
             The range of bytes to read.
 
@@ -164,7 +165,7 @@ async def get(
             The read bytes, or None if the key does not exist.
         """
         if prototype is None:
-            prototype = default_buffer_prototype()
+            prototype = self.store._get_default_buffer_class()
         return await self.store.get(self.path, prototype=prototype, byte_range=byte_range)
 
     async def set(self, value: Buffer) -> None: