Use config to select implementation (#1982)

* make codec pipeline implementation configurable

* add test_config_codec_pipeline_class_in_env

* make codec implementation configurable

* remove snake case support for class names in config

* use registry for codec pipeline config

* typing

* load codec pipeline from entrypoints

* test if configured codec implementation and codec pipeline is used

* make ndbuffer implementation configurable

* fix circular import

* change class method calls on NDBuffer to use get_ndbuffer_class()

* make buffer implementation configurable

* format

* fix tests

* ignore mypy in tests

* add test to lazy load (nd)buffer from entrypoint

* better assertion message

* fix merge

* fix merge

* formatting

* fix mypy

* fix ruff formatting

* fix merge

* fix mypy

* use numpy_buffer_prototype for reading shard index

* rename buffer and entrypoint test-classes

* document interaction registry and config

* change config prefix from zarr_python to zarr

* use fully_qualified_name for implementation config

* refactor registry dicts

* fix default_buffer_prototype access in tests

* allow multiple implementations per entry_point

* add tests for multiple implementations per entry_point

* fix DeprecationWarning: SelectableGroups in registry.py

* fix DeprecationWarning: EntryPoints list interface in registry.py

* clarify  _collect_entrypoints docstring

Co-authored-by: Norman Rzepka <[email protected]>

---------

Co-authored-by: Norman Rzepka <[email protected]>

Author: brokkoli71

src/zarr/abc/codec.py

@@ -17,6 +17,7 @@
     from typing_extensions import Self
 
     from zarr.array_spec import ArraySpec
+    from zarr.common import JSON
     from zarr.indexing import SelectorTuple
 
 CodecInput = TypeVar("CodecInput", bound=NDBuffer | Buffer)
@@ -254,7 +255,7 @@ def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self:
 
     @classmethod
     @abstractmethod
-    def from_list(cls, codecs: list[Codec]) -> Self:
+    def from_list(cls, codecs: Iterable[Codec]) -> Self:
         """Creates a codec pipeline from a list of codecs.
 
         Parameters
@@ -390,6 +391,15 @@ async def write(
         """
         ...
 
+    @classmethod
+    def from_dict(cls, data: Iterable[JSON | Codec]) -> Self:
+        """
+        Create an instance of the model from a dictionary
+        """
+        ...
+
+        return cls(**data)
+
 
 async def batching_helper(
     func: Callable[[CodecInput, ArraySpec], Awaitable[CodecOutput | None]],

src/zarr/array.py

@@ -25,7 +25,6 @@
 from zarr.chunk_key_encodings import ChunkKeyEncoding, DefaultChunkKeyEncoding, V2ChunkKeyEncoding
 from zarr.codecs import BytesCodec
 from zarr.codecs._v2 import V2Compressor, V2Filters
-from zarr.codecs.pipeline import BatchedCodecPipeline
 from zarr.common import (
     JSON,
     ZARR_JSON,
@@ -61,6 +60,7 @@
     pop_fields,
 )
 from zarr.metadata import ArrayMetadata, ArrayV2Metadata, ArrayV3Metadata
+from zarr.registry import get_pipeline_class
 from zarr.store import StoreLike, StorePath, make_store_path
 from zarr.store.core import (
     ensure_no_existing_node,
@@ -79,11 +79,11 @@ def parse_array_metadata(data: Any) -> ArrayV2Metadata | ArrayV3Metadata:
     raise TypeError
 
 
-def create_codec_pipeline(metadata: ArrayV2Metadata | ArrayV3Metadata) -> BatchedCodecPipeline:
+def create_codec_pipeline(metadata: ArrayV2Metadata | ArrayV3Metadata) -> CodecPipeline:
     if isinstance(metadata, ArrayV3Metadata):
-        return BatchedCodecPipeline.from_list(metadata.codecs)
+        return get_pipeline_class().from_list(metadata.codecs)
     elif isinstance(metadata, ArrayV2Metadata):
-        return BatchedCodecPipeline.from_list(
+        return get_pipeline_class().from_list(
             [V2Filters(metadata.filters or []), V2Compressor(metadata.compressor)]
         )
     else:
@@ -483,8 +483,13 @@ async def _get_selection(
         return out_buffer.as_ndarray_like()
 
     async def getitem(
-        self, selection: BasicSelection, *, prototype: BufferPrototype = default_buffer_prototype
+        self,
+        selection: BasicSelection,
+        *,
+        prototype: BufferPrototype | None = None,
     ) -> NDArrayLike:
+        if prototype is None:
+            prototype = default_buffer_prototype()
         indexer = BasicIndexer(
             selection,
             shape=self.metadata.shape,
@@ -493,7 +498,7 @@ async def getitem(
         return await self._get_selection(indexer, prototype=prototype)
 
     async def _save_metadata(self, metadata: ArrayMetadata) -> None:
-        to_save = metadata.to_buffer_dict()
+        to_save = metadata.to_buffer_dict(default_buffer_prototype())
         awaitables = [set_or_delete(self.store_path / key, value) for key, value in to_save.items()]
         await gather(*awaitables)
 
@@ -545,8 +550,10 @@ async def setitem(
         self,
         selection: BasicSelection,
         value: npt.ArrayLike,
-        prototype: BufferPrototype = default_buffer_prototype,
+        prototype: BufferPrototype | None = None,
     ) -> None:
+        if prototype is None:
+            prototype = default_buffer_prototype()
         indexer = BasicIndexer(
             selection,
             shape=self.metadata.shape,
@@ -1001,7 +1008,7 @@ def get_basic_selection(
         selection: BasicSelection = Ellipsis,
         *,
         out: NDBuffer | None = None,
-        prototype: BufferPrototype = default_buffer_prototype,
+        prototype: BufferPrototype | None = None,
         fields: Fields | None = None,
     ) -> NDArrayLike:
         """Retrieve data for an item or region of the array.
@@ -1108,6 +1115,8 @@ def get_basic_selection(
 
         """
 
+        if prototype is None:
+            prototype = default_buffer_prototype()
         return sync(
             self._async_array._get_selection(
                 BasicIndexer(selection, self.shape, self.metadata.chunk_grid),
@@ -1123,7 +1132,7 @@ def set_basic_selection(
         value: npt.ArrayLike,
         *,
         fields: Fields | None = None,
-        prototype: BufferPrototype = default_buffer_prototype,
+        prototype: BufferPrototype | None = None,
     ) -> None:
         """Modify data for an item or region of the array.
 
@@ -1207,6 +1216,8 @@ def set_basic_selection(
         vindex, oindex, blocks, __getitem__, __setitem__
 
         """
+        if prototype is None:
+            prototype = default_buffer_prototype()
         indexer = BasicIndexer(selection, self.shape, self.metadata.chunk_grid)
         sync(self._async_array._set_selection(indexer, value, fields=fields, prototype=prototype))
 
@@ -1216,7 +1227,7 @@ def get_orthogonal_selection(
         *,
         out: NDBuffer | None = None,
         fields: Fields | None = None,
-        prototype: BufferPrototype = default_buffer_prototype,
+        prototype: BufferPrototype | None = None,
     ) -> NDArrayLike:
         """Retrieve data by making a selection for each dimension of the array. For
         example, if an array has 2 dimensions, allows selecting specific rows and/or
@@ -1325,6 +1336,8 @@ def get_orthogonal_selection(
         vindex, oindex, blocks, __getitem__, __setitem__
 
         """
+        if prototype is None:
+            prototype = default_buffer_prototype()
         indexer = OrthogonalIndexer(selection, self.shape, self.metadata.chunk_grid)
         return sync(
             self._async_array._get_selection(
@@ -1338,7 +1351,7 @@ def set_orthogonal_selection(
         value: npt.ArrayLike,
         *,
         fields: Fields | None = None,
-        prototype: BufferPrototype = default_buffer_prototype,
+        prototype: BufferPrototype | None = None,
     ) -> None:
         """Modify data via a selection for each dimension of the array.
 
@@ -1435,6 +1448,8 @@ def set_orthogonal_selection(
         vindex, oindex, blocks, __getitem__, __setitem__
 
         """
+        if prototype is None:
+            prototype = default_buffer_prototype()
         indexer = OrthogonalIndexer(selection, self.shape, self.metadata.chunk_grid)
         return sync(
             self._async_array._set_selection(indexer, value, fields=fields, prototype=prototype)
@@ -1446,7 +1461,7 @@ def get_mask_selection(
         *,
         out: NDBuffer | None = None,
         fields: Fields | None = None,
-        prototype: BufferPrototype = default_buffer_prototype,
+        prototype: BufferPrototype | None = None,
     ) -> NDArrayLike:
         """Retrieve a selection of individual items, by providing a Boolean array of the
         same shape as the array against which the selection is being made, where True
@@ -1513,6 +1528,8 @@ def get_mask_selection(
         vindex, oindex, blocks, __getitem__, __setitem__
         """
 
+        if prototype is None:
+            prototype = default_buffer_prototype()
         indexer = MaskIndexer(mask, self.shape, self.metadata.chunk_grid)
         return sync(
             self._async_array._get_selection(
@@ -1526,7 +1543,7 @@ def set_mask_selection(
         value: npt.ArrayLike,
         *,
         fields: Fields | None = None,
-        prototype: BufferPrototype = default_buffer_prototype,
+        prototype: BufferPrototype | None = None,
     ) -> None:
         """Modify a selection of individual items, by providing a Boolean array of the
         same shape as the array against which the selection is being made, where True
@@ -1593,6 +1610,8 @@ def set_mask_selection(
         vindex, oindex, blocks, __getitem__, __setitem__
 
         """
+        if prototype is None:
+            prototype = default_buffer_prototype()
         indexer = MaskIndexer(mask, self.shape, self.metadata.chunk_grid)
         sync(self._async_array._set_selection(indexer, value, fields=fields, prototype=prototype))
 
@@ -1602,7 +1621,7 @@ def get_coordinate_selection(
         *,
         out: NDBuffer | None = None,
         fields: Fields | None = None,
-        prototype: BufferPrototype = default_buffer_prototype,
+        prototype: BufferPrototype | None = None,
     ) -> NDArrayLike:
         """Retrieve a selection of individual items, by providing the indices
         (coordinates) for each selected item.
@@ -1671,6 +1690,8 @@ def get_coordinate_selection(
         vindex, oindex, blocks, __getitem__, __setitem__
 
         """
+        if prototype is None:
+            prototype = default_buffer_prototype()
         indexer = CoordinateIndexer(selection, self.shape, self.metadata.chunk_grid)
         out_array = sync(
             self._async_array._get_selection(
@@ -1689,7 +1710,7 @@ def set_coordinate_selection(
         value: npt.ArrayLike,
         *,
         fields: Fields | None = None,
-        prototype: BufferPrototype = default_buffer_prototype,
+        prototype: BufferPrototype | None = None,
     ) -> None:
         """Modify a selection of individual items, by providing the indices (coordinates)
         for each item to be modified.
@@ -1753,6 +1774,8 @@ def set_coordinate_selection(
         vindex, oindex, blocks, __getitem__, __setitem__
 
         """
+        if prototype is None:
+            prototype = default_buffer_prototype()
         # setup indexer
         indexer = CoordinateIndexer(selection, self.shape, self.metadata.chunk_grid)
 
@@ -1776,7 +1799,7 @@ def get_block_selection(
         *,
         out: NDBuffer | None = None,
         fields: Fields | None = None,
-        prototype: BufferPrototype = default_buffer_prototype,
+        prototype: BufferPrototype | None = None,
     ) -> NDArrayLike:
         """Retrieve a selection of individual items, by providing the indices
         (coordinates) for each selected item.
@@ -1859,6 +1882,8 @@ def get_block_selection(
         vindex, oindex, blocks, __getitem__, __setitem__
 
         """
+        if prototype is None:
+            prototype = default_buffer_prototype()
         indexer = BlockIndexer(selection, self.shape, self.metadata.chunk_grid)
         return sync(
             self._async_array._get_selection(
@@ -1872,7 +1897,7 @@ def set_block_selection(
         value: npt.ArrayLike,
         *,
         fields: Fields | None = None,
-        prototype: BufferPrototype = default_buffer_prototype,
+        prototype: BufferPrototype | None = None,
     ) -> None:
         """Modify a selection of individual blocks, by providing the chunk indices
         (coordinates) for each block to be modified.
@@ -1950,6 +1975,8 @@ def set_block_selection(
         vindex, oindex, blocks, __getitem__, __setitem__
 
         """
+        if prototype is None:
+            prototype = default_buffer_prototype()
         indexer = BlockIndexer(selection, self.shape, self.metadata.chunk_grid)
         sync(self._async_array._set_selection(indexer, value, fields=fields, prototype=prototype))
 

src/zarr/buffer.py

@@ -16,6 +16,12 @@
 import numpy.typing as npt
 
 from zarr.common import ChunkCoords
+from zarr.registry import (
+    get_buffer_class,
+    get_ndbuffer_class,
+    register_buffer,
+    register_ndbuffer,
+)
 
 if TYPE_CHECKING:
     from typing_extensions import Self
@@ -479,4 +485,14 @@ class BufferPrototype(NamedTuple):
 
 
 # The default buffer prototype used throughout the Zarr codebase.
-default_buffer_prototype = BufferPrototype(buffer=Buffer, nd_buffer=NDBuffer)
+def default_buffer_prototype() -> BufferPrototype:
+    return BufferPrototype(buffer=get_buffer_class(), nd_buffer=get_ndbuffer_class())
+
+
+# The numpy prototype used for E.g. when reading the shard index
+def numpy_buffer_prototype() -> BufferPrototype:
+    return BufferPrototype(buffer=Buffer, nd_buffer=NDBuffer)
+
+
+register_buffer(Buffer)
+register_ndbuffer(NDBuffer)

src/zarr/codecs/_v2.py

@@ -7,8 +7,9 @@
 
 from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec
 from zarr.array_spec import ArraySpec
-from zarr.buffer import Buffer, NDBuffer
+from zarr.buffer import Buffer, NDBuffer, default_buffer_prototype
 from zarr.common import JSON, to_thread
+from zarr.registry import get_ndbuffer_class
 
 
 @dataclass(frozen=True)
@@ -34,7 +35,7 @@ async def _decode_single(
         if str(chunk_numpy_array.dtype) != chunk_spec.dtype:
             chunk_numpy_array = chunk_numpy_array.view(chunk_spec.dtype)
 
-        return NDBuffer.from_numpy_array(chunk_numpy_array)
+        return get_ndbuffer_class().from_numpy_array(chunk_numpy_array)
 
     async def _encode_single(
         self,
@@ -55,7 +56,7 @@ async def _encode_single(
         else:
             encoded_chunk_bytes = ensure_bytes(chunk_numpy_array)
 
-        return Buffer.from_bytes(encoded_chunk_bytes)
+        return default_buffer_prototype().buffer.from_bytes(encoded_chunk_bytes)
 
     def compute_encoded_size(self, _input_byte_length: int, _chunk_spec: ArraySpec) -> int:
         raise NotImplementedError
@@ -86,7 +87,7 @@ async def _decode_single(
                 order=chunk_spec.order,
             )
 
-        return NDBuffer.from_ndarray_like(chunk_ndarray)
+        return get_ndbuffer_class().from_ndarray_like(chunk_ndarray)
 
     async def _encode_single(
         self,
@@ -99,7 +100,7 @@ async def _encode_single(
             filter = numcodecs.get_codec(filter_metadata)
             chunk_ndarray = await to_thread(filter.encode, chunk_ndarray)
 
-        return NDBuffer.from_ndarray_like(chunk_ndarray)
+        return get_ndbuffer_class().from_ndarray_like(chunk_ndarray)
 
     def compute_encoded_size(self, _input_byte_length: int, _chunk_spec: ArraySpec) -> int:
         raise NotImplementedError

src/zarr/codecs/blosc.py

@@ -11,8 +11,8 @@
 from zarr.abc.codec import BytesBytesCodec
 from zarr.array_spec import ArraySpec
 from zarr.buffer import Buffer, as_numpy_array_wrapper
-from zarr.codecs.registry import register_codec
 from zarr.common import JSON, parse_enum, parse_named_configuration, to_thread
+from zarr.registry import register_codec
 
 if TYPE_CHECKING:
     from typing_extensions import Self

src/zarr/codecs/bytes.py

@@ -10,8 +10,8 @@
 from zarr.abc.codec import ArrayBytesCodec
 from zarr.array_spec import ArraySpec
 from zarr.buffer import Buffer, NDArrayLike, NDBuffer
-from zarr.codecs.registry import register_codec
 from zarr.common import JSON, parse_enum, parse_named_configuration
+from zarr.registry import register_codec
 
 if TYPE_CHECKING:
     from typing_extensions import Self

src/zarr/codecs/crc32c_.py

@@ -9,8 +9,8 @@
 from zarr.abc.codec import BytesBytesCodec
 from zarr.array_spec import ArraySpec
 from zarr.buffer import Buffer
-from zarr.codecs.registry import register_codec
 from zarr.common import JSON, parse_named_configuration
+from zarr.registry import register_codec
 
 if TYPE_CHECKING:
     from typing_extensions import Self

src/zarr/codecs/gzip.py

@@ -8,8 +8,8 @@
 from zarr.abc.codec import BytesBytesCodec
 from zarr.array_spec import ArraySpec
 from zarr.buffer import Buffer, as_numpy_array_wrapper
-from zarr.codecs.registry import register_codec
 from zarr.common import JSON, parse_named_configuration, to_thread
+from zarr.registry import register_codec
 
 if TYPE_CHECKING:
     from typing_extensions import Self