remove as much as possible default_buffer_prototype() invocation

Author: d-v-b

src/zarr/abc/store.py

@@ -8,17 +8,18 @@
 from typing import TYPE_CHECKING, Literal, Protocol, runtime_checkable
 
 from zarr.core.buffer import Buffer, BufferPrototype
-from zarr.core.buffer.core import default_buffer_prototype
 from zarr.core.sync import sync
+from zarr.registry import get_buffer_class
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, AsyncIterator, Iterable
     from types import TracebackType
     from typing import Any, Self, TypeAlias
 
-__all__ = ["BufferLike", "ByteGetter", "ByteSetter", "Store", "set_or_delete"]
+__all__ = ["BufferClassLike", "ByteGetter", "ByteSetter", "Store", "set_or_delete"]
 
-BufferLike = type[Buffer] | BufferPrototype
+BufferClassLike = type[Buffer] | BufferPrototype
+"""An object that is or contains a Buffer class"""
 
 
 @dataclass
@@ -189,13 +190,13 @@ def _get_default_buffer_class(self) -> type[Buffer]:
         """
         Get the default buffer class.
         """
-        return default_buffer_prototype().buffer
+        return get_buffer_class()
 
     @abstractmethod
     async def get(
         self,
         key: str,
-        prototype: BufferLike | None = None,
+        prototype: BufferClassLike | None = None,
         byte_range: ByteRequest | None = None,
     ) -> Buffer | None:
         """Retrieve the value associated with a given key.
@@ -225,7 +226,7 @@ async def _get_bytes(
         self,
         key: str,
         *,
-        prototype: BufferLike | None = None,
+        prototype: BufferClassLike | None = None,
         byte_range: ByteRequest | None = None,
     ) -> bytes:
         """
@@ -268,7 +269,7 @@ async def _get_bytes(
         --------
         >>> store = await MemoryStore.open()
         >>> await store.set("data", Buffer.from_bytes(b"hello world"))
-        >>> data = await store._get_bytes("data", prototype=default_buffer_prototype())
+        >>> data = await store._get_bytes("data")
         >>> print(data)
         b'hello world'
         """
@@ -281,7 +282,7 @@ def _get_bytes_sync(
         self,
         key: str = "",
         *,
-        prototype: BufferLike | None = None,
+        prototype: BufferClassLike | None = None,
         byte_range: ByteRequest | None = None,
     ) -> bytes:
         """
@@ -329,7 +330,7 @@ def _get_bytes_sync(
         --------
         >>> store = MemoryStore()
         >>> await store.set("data", Buffer.from_bytes(b"hello world"))
-        >>> data = store._get_bytes_sync("data", prototype=default_buffer_prototype())
+        >>> data = store._get_bytes_sync("data")
         >>> print(data)
         b'hello world'
         """
@@ -340,7 +341,7 @@ async def _get_json(
         self,
         key: str,
         *,
-        prototype: BufferLike | None = None,
+        prototype: BufferClassLike | None = None,
         byte_range: ByteRequest | None = None,
     ) -> Any:
         """
@@ -387,7 +388,7 @@ async def _get_json(
         >>> store = await MemoryStore.open()
         >>> metadata = {"zarr_format": 3, "node_type": "array"}
         >>> await store.set("zarr.json", Buffer.from_bytes(json.dumps(metadata).encode()))
-        >>> data = await store._get_json("zarr.json", prototype=default_buffer_prototype())
+        >>> data = await store._get_json("zarr.json")
         >>> print(data)
         {'zarr_format': 3, 'node_type': 'array'}
         """
@@ -398,7 +399,7 @@ def _get_json_sync(
         self,
         key: str = "",
         *,
-        prototype: BufferLike | None = None,
+        prototype: BufferClassLike | None = None,
         byte_range: ByteRequest | None = None,
     ) -> Any:
         """
@@ -451,7 +452,7 @@ def _get_json_sync(
         >>> store = MemoryStore()
         >>> metadata = {"zarr_format": 3, "node_type": "array"}
         >>> store.set("zarr.json", Buffer.from_bytes(json.dumps(metadata).encode()))
-        >>> data = store._get_json_sync("zarr.json", prototype=default_buffer_prototype())
+        >>> data = store._get_json_sync("zarr.json")
         >>> print(data)
         {'zarr_format': 3, 'node_type': 'array'}
         """
@@ -461,7 +462,7 @@ def _get_json_sync(
     @abstractmethod
     async def get_partial_values(
         self,
-        prototype: BufferLike | None,
+        prototype: BufferClassLike | None,
         key_ranges: Iterable[tuple[str, ByteRequest | None]],
     ) -> list[Buffer | None]:
         """Retrieve possibly partial values from given key_ranges.
@@ -645,7 +646,7 @@ def close(self) -> None:
         self._is_open = False
 
     async def _get_many(
-        self, requests: Iterable[tuple[str, BufferLike | None, ByteRequest | None]]
+        self, requests: Iterable[tuple[str, BufferClassLike | None, ByteRequest | None]]
     ) -> AsyncGenerator[tuple[str, Buffer | None], None]:
         """
         Retrieve a collection of objects from storage. In general this method does not guarantee
@@ -676,10 +677,8 @@ async def getsize(self, key: str) -> int:
         # Note to implementers: this default implementation is very inefficient since
         # it requires reading the entire object. Many systems will have ways to get the
         # size of an object without reading it.
-        # avoid circular import
-        from zarr.core.buffer.core import default_buffer_prototype
 
-        value = await self.get(key, prototype=default_buffer_prototype())
+        value = await self.get(key)
         if value is None:
             raise FileNotFoundError(key)
         return len(value)

src/zarr/core/common.py

@@ -23,12 +23,15 @@
 
 from typing_extensions import ReadOnly
 
+from zarr.core.buffer import Buffer, BufferPrototype
 from zarr.core.config import config as zarr_config
 from zarr.errors import ZarrRuntimeWarning
 
 if TYPE_CHECKING:
     from collections.abc import Awaitable, Callable, Iterator
 
+    from zarr.abc.store import BufferClassLike
+
 
 ZARR_JSON = "zarr.json"
 ZARRAY_JSON = ".zarray"
@@ -246,3 +249,17 @@ def _warn_order_kwarg() -> None:
 def _default_zarr_format() -> ZarrFormat:
     """Return the default zarr_version"""
     return cast("ZarrFormat", int(zarr_config.get("default_zarr_format", 3)))
+
+
+def parse_bufferclasslike(obj: BufferClassLike | None) -> type[Buffer]:
+    """
+    Take an optional BufferClassLike and return a Buffer class
+    """
+    # Avoid a circular import. Temporary fix until we re-organize modules appropriately.
+    from zarr.registry import get_buffer_class
+
+    if obj is None:
+        return get_buffer_class()
+    if isinstance(obj, BufferPrototype):
+        return obj.buffer
+    return obj

src/zarr/core/group.py

@@ -17,7 +17,7 @@
 
 import zarr.api.asynchronous as async_api
 from zarr.abc.metadata import Metadata
-from zarr.abc.store import Store, set_or_delete
+from zarr.abc.store import BufferClassLike, Store, set_or_delete
 from zarr.core._info import GroupInfo
 from zarr.core.array import (
     DEFAULT_FILL_VALUE,
@@ -32,7 +32,6 @@
     create_array,
 )
 from zarr.core.attributes import Attributes
-from zarr.core.buffer import default_buffer_prototype
 from zarr.core.common import (
     JSON,
     ZARR_JSON,
@@ -44,6 +43,7 @@
     NodeType,
     ShapeLike,
     ZarrFormat,
+    parse_bufferclasslike,
     parse_shapelike,
 )
 from zarr.core.config import config
@@ -75,7 +75,7 @@
     from typing import Any
 
     from zarr.core.array_spec import ArrayConfigLike
-    from zarr.core.buffer import Buffer, BufferPrototype
+    from zarr.core.buffer import Buffer
     from zarr.core.chunk_key_encodings import ChunkKeyEncodingLike
     from zarr.core.common import MemoryOrder
     from zarr.core.dtype import ZDTypeLike
@@ -356,20 +356,25 @@ class GroupMetadata(Metadata):
     consolidated_metadata: ConsolidatedMetadata | None = None
     node_type: Literal["group"] = field(default="group", init=False)
 
-    def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]:
+    def to_buffer_dict(self, prototype: BufferClassLike | None = None) -> dict[str, Buffer]:
+        """
+        Convert the metadata document to a dict with string keys and `Buffer` values.
+        """
+        buffer_cls = parse_bufferclasslike(prototype)
+
         json_indent = config.get("json_indent")
         if self.zarr_format == 3:
             return {
-                ZARR_JSON: prototype.buffer.from_bytes(
+                ZARR_JSON: buffer_cls.from_bytes(
                     json.dumps(self.to_dict(), indent=json_indent, allow_nan=True).encode()
                 )
             }
         else:
             items = {
-                ZGROUP_JSON: prototype.buffer.from_bytes(
+                ZGROUP_JSON: buffer_cls.from_bytes(
                     json.dumps({"zarr_format": self.zarr_format}, indent=json_indent).encode()
                 ),
-                ZATTRS_JSON: prototype.buffer.from_bytes(
+                ZATTRS_JSON: buffer_cls.from_bytes(
                     json.dumps(self.attributes, indent=json_indent, allow_nan=True).encode()
                 ),
             }
@@ -396,7 +401,7 @@ def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]:
                             },
                         }
 
-                items[ZMETADATA_V2_JSON] = prototype.buffer.from_bytes(
+                items[ZMETADATA_V2_JSON] = buffer_cls.from_bytes(
                     json.dumps(
                         {"metadata": d, "zarr_consolidated_format": 1}, allow_nan=True
                     ).encode()
@@ -2029,7 +2034,7 @@ async def update_attributes_async(self, new_attributes: dict[str, Any]) -> Group
         new_metadata = replace(self.metadata, attributes=new_attributes)
 
         # Write new metadata
-        to_save = new_metadata.to_buffer_dict(default_buffer_prototype())
+        to_save = new_metadata.to_buffer_dict()
         awaitables = [set_or_delete(self.store_path / key, value) for key, value in to_save.items()]
         await asyncio.gather(*awaitables)
 
@@ -3615,9 +3620,7 @@ async def _read_metadata_v3(store: Store, path: str) -> ArrayV3Metadata | GroupM
     document stored at store_path.path / zarr.json. If no such document is found, raise a
     FileNotFoundError.
     """
-    zarr_json_bytes = await store.get(
-        _join_paths([path, ZARR_JSON]), prototype=default_buffer_prototype()
-    )
+    zarr_json_bytes = await store.get(_join_paths([path, ZARR_JSON]))
     if zarr_json_bytes is None:
         raise FileNotFoundError(path)
     else:
@@ -3634,9 +3637,9 @@ async def _read_metadata_v2(store: Store, path: str) -> ArrayV2Metadata | GroupM
     # TODO: consider first fetching array metadata, and only fetching group metadata when we don't
     # find an array
     zarray_bytes, zgroup_bytes, zattrs_bytes = await asyncio.gather(
-        store.get(_join_paths([path, ZARRAY_JSON]), prototype=default_buffer_prototype()),
-        store.get(_join_paths([path, ZGROUP_JSON]), prototype=default_buffer_prototype()),
-        store.get(_join_paths([path, ZATTRS_JSON]), prototype=default_buffer_prototype()),
+        store.get(_join_paths([path, ZARRAY_JSON])),
+        store.get(_join_paths([path, ZGROUP_JSON])),
+        store.get(_join_paths([path, ZATTRS_JSON])),
     )
 
     if zattrs_bytes is None:
@@ -3850,7 +3853,7 @@ def _persist_metadata(
     Prepare to save a metadata document to storage, returning a tuple of coroutines that must be awaited.
     """
 
-    to_save = metadata.to_buffer_dict(default_buffer_prototype())
+    to_save = metadata.to_buffer_dict()
     return tuple(
         _set_return_key(store=store, key=_join_paths([path, key]), value=value, semaphore=semaphore)
         for key, value in to_save.items()

src/zarr/core/metadata/io.py

@@ -4,7 +4,6 @@
 from typing import TYPE_CHECKING
 
 from zarr.abc.store import set_or_delete
-from zarr.core.buffer.core import default_buffer_prototype
 from zarr.errors import ContainsArrayError
 from zarr.storage._common import StorePath, ensure_no_existing_node
 
@@ -51,7 +50,7 @@ async def save_metadata(
     ------
     ValueError
     """
-    to_save = metadata.to_buffer_dict(default_buffer_prototype())
+    to_save = metadata.to_buffer_dict()
     set_awaitables = [set_or_delete(store_path / key, value) for key, value in to_save.items()]
 
     if ensure_parents:
@@ -71,9 +70,7 @@ async def save_metadata(
             set_awaitables.extend(
                 [
                     (parent_store_path / key).set_if_not_exists(value)
-                    for key, value in parent_metadata.to_buffer_dict(
-                        default_buffer_prototype()
-                    ).items()
+                    for key, value in parent_metadata.to_buffer_dict().items()
                 ]
             )
 

src/zarr/core/metadata/v2.py

@@ -18,6 +18,7 @@
 
     import numpy.typing as npt
 
+    from zarr.abc.store import BufferClassLike
     from zarr.core.buffer import Buffer, BufferPrototype
     from zarr.core.dtype.wrapper import (
         TBaseDType,
@@ -39,6 +40,7 @@
     ZARRAY_JSON,
     ZATTRS_JSON,
     MemoryOrder,
+    parse_bufferclasslike,
     parse_shapelike,
 )
 from zarr.core.config import config, parse_indexing_order
@@ -125,15 +127,16 @@ def chunk_grid(self) -> RegularChunkGrid:
     def shards(self) -> tuple[int, ...] | None:
         return None
 
-    def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]:
+    def to_buffer_dict(self, prototype: BufferClassLike | None = None) -> dict[str, Buffer]:
+        buffer_cls = parse_bufferclasslike(prototype)
         zarray_dict = self.to_dict()
         zattrs_dict = zarray_dict.pop("attributes", {})
         json_indent = config.get("json_indent")
         return {
-            ZARRAY_JSON: prototype.buffer.from_bytes(
+            ZARRAY_JSON: buffer_cls.from_bytes(
                 json.dumps(zarray_dict, indent=json_indent, allow_nan=True).encode()
             ),
-            ZATTRS_JSON: prototype.buffer.from_bytes(
+            ZATTRS_JSON: buffer_cls.from_bytes(
                 json.dumps(zattrs_dict, indent=json_indent, allow_nan=True).encode()
             ),
         }

src/zarr/core/metadata/v3.py

@@ -11,6 +11,7 @@
 if TYPE_CHECKING:
     from typing import Self
 
+    from zarr.abc.store import BufferClassLike
     from zarr.core.buffer import Buffer, BufferPrototype
     from zarr.core.chunk_grids import ChunkGrid
     from zarr.core.common import JSON
@@ -35,6 +36,7 @@
     ZARR_JSON,
     DimensionNames,
     NamedConfig,
+    parse_bufferclasslike,
     parse_named_configuration,
     parse_shapelike,
 )
@@ -345,11 +347,12 @@ def get_chunk_spec(
     def encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str:
         return self.chunk_key_encoding.encode_chunk_key(chunk_coords)
 
-    def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]:
+    def to_buffer_dict(self, prototype: BufferClassLike | None = None) -> dict[str, Buffer]:
+        buffer_cls = parse_bufferclasslike(prototype)
         json_indent = config.get("json_indent")
         d = self.to_dict()
         return {
-            ZARR_JSON: prototype.buffer.from_bytes(
+            ZARR_JSON: buffer_cls.from_bytes(
                 json.dumps(d, allow_nan=True, indent=json_indent).encode()
             )
         }