Wip implementation of v2 / v3 codec behavior

Author: d-v-b

src/zarr/abc/codec.py

@@ -2,6 +2,7 @@
 
 from abc import abstractmethod
 from collections.abc import Mapping
+from dataclasses import dataclass
 from typing import (
     TYPE_CHECKING,
     ClassVar,
@@ -525,3 +526,5 @@ def get_config(self) -> CodecJSON_V2[str]: ...
 
     @classmethod
     def from_config(cls, config: CodecJSON_V2[str]) -> Self: ...
+
+

src/zarr/abc/store.py

@@ -6,10 +6,6 @@
 from itertools import starmap
 from typing import TYPE_CHECKING, Protocol, runtime_checkable
 
-from zarr.core.buffer.core import default_buffer_prototype
-from zarr.core.common import concurrent_map
-from zarr.core.config import config
-
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, AsyncIterator, Iterable
     from types import TracebackType
@@ -435,6 +431,7 @@ async def getsize(self, key: str) -> int:
         FileNotFoundError
             When the given key does not exist in the store.
         """
+        from zarr.core.buffer.core import default_buffer_prototype
         # 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.
@@ -476,6 +473,8 @@ async def getsize_prefix(self, prefix: str) -> int:
         # on to getting sizes. Ideally we would overlap those two, which should
         # improve tail latency and might reduce memory pressure (since not all keys
         # would be in memory at once).
+        from zarr.core.common import concurrent_map
+        from zarr.core.config import config
         keys = [(x,) async for x in self.list_prefix(prefix)]
         limit = config.get("async.concurrency")
         sizes = await concurrent_map(keys, self.getsize, limit=limit)

src/zarr/codecs/_v2.py

@@ -2,12 +2,15 @@
 
 import asyncio
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, TypeGuard
+from typing import TYPE_CHECKING, Self, TypeGuard, overload
+from typing_extensions import Literal
 
 import numpy as np
 from numcodecs.compat import ensure_bytes, ensure_ndarray_like
 
-from zarr.abc.codec import ArrayBytesCodec, Numcodec
+from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BaseCodec, BytesBytesCodec, CodecJSON, CodecJSON_V2, Numcodec
+from zarr.core.buffer.core import BufferPrototype
+from zarr.core.common import BaseConfig, NamedConfig, ZarrFormat
 from zarr.registry import get_ndbuffer_class
 
 if TYPE_CHECKING:
@@ -133,3 +136,100 @@ async def _encode_single(
 
     def compute_encoded_size(self, _input_byte_length: int, _chunk_spec: ArraySpec) -> int:
         raise NotImplementedError
+
+
+@dataclass(frozen=True, kw_only=True)
+class NumcodecsWrapper:
+    codec: Numcodec
+
+    @overload
+    def to_json(self, zarr_format: Literal[2]) -> CodecJSON_V2[str]: ...
+    @overload
+    def to_json(self, zarr_format: Literal[3]) -> NamedConfig[str, BaseConfig]: ...
+
+    def to_json(self, zarr_format: ZarrFormat) -> CodecJSON_V2[str] | NamedConfig[str, BaseConfig]:
+        if zarr_format == 2:
+            return self.codec.get_config()
+        elif zarr_format == 3:
+            config = self.codec.get_config()
+            config_no_id = {k: v for k, v in config.items() if k != "id"}
+            return {"name": config["id"], "configuration": config_no_id}
+        raise ValueError(f"Unsupported zarr format: {zarr_format}")  # pragma: no cover
+
+    @classmethod
+    def _from_json_v2(cls, data: CodecJSON) -> Self:
+        raise NotADirectoryError(
+            "This class does not support creating instances from JSON data for Zarr format 2."
+            )
+
+    @classmethod
+    def _from_json_v3(cls, data: CodecJSON) -> Self:
+        raise NotImplementedError(
+            "This class does not support creating instances from JSON data for Zarr format 3."
+        )
+
+    def compute_encoded_size(self, input_byte_length: int, chunk_spec: ArraySpec) -> int:
+        raise NotImplementedError
+
+    def to_array_array(self) -> NumcodecsArrayArrayCodec:
+        """
+        Use the ``_codec`` attribute to create a NumcodecsArrayArrayCodec.
+        """
+        return NumcodecsArrayArrayCodec(codec=self.codec)
+
+    def to_bytes_bytes(self) -> NumcodecsBytesBytesCodec:
+        """
+        Use the ``_codec`` attribute to create a NumcodecsBytesBytesCodec.
+        """
+        return NumcodecsBytesBytesCodec(codec=self.codec)
+
+    def to_array_bytes(self) -> NumcodecsArrayBytesCodec:
+        """
+        Use the ``_codec`` attribute to create a NumcodecsArrayBytesCodec.
+        """
+        return NumcodecsArrayBytesCodec(codec=self.codec)
+
+
+class NumcodecsBytesBytesCodec(NumcodecsWrapper, BytesBytesCodec):
+    async def _decode_single(self, chunk_data: Buffer, chunk_spec: ArraySpec) -> Buffer:
+        return await asyncio.to_thread(
+            as_numpy_array_wrapper,
+            self.codec.decode,
+            chunk_data,
+            chunk_spec.prototype,
+        )
+
+    def _encode(self, chunk_bytes: Buffer, prototype: BufferPrototype) -> Buffer:
+        encoded = self.codec.encode(chunk_bytes.as_array_like())
+        if isinstance(encoded, np.ndarray):  # Required for checksum codecs
+            return prototype.buffer.from_bytes(encoded.tobytes())
+        return prototype.buffer.from_bytes(encoded)
+
+    async def _encode_single(self, chunk_data: Buffer, chunk_spec: ArraySpec) -> Buffer:
+        return await asyncio.to_thread(self._encode, chunk_data, chunk_spec.prototype)
+
+
+@dataclass(kw_only=True, frozen=True)
+class NumcodecsArrayArrayCodec(NumcodecsWrapper, ArrayArrayCodec):
+    async def _decode_single(self, chunk_data: NDBuffer, chunk_spec: ArraySpec) -> NDBuffer:
+        chunk_ndarray = chunk_data.as_ndarray_like()
+        out = await asyncio.to_thread(self.codec.decode, chunk_ndarray)
+        return chunk_spec.prototype.nd_buffer.from_ndarray_like(out.reshape(chunk_spec.shape))  # type: ignore[union-attr]
+
+    async def _encode_single(self, chunk_data: NDBuffer, chunk_spec: ArraySpec) -> NDBuffer:
+        chunk_ndarray = chunk_data.as_ndarray_like()
+        out = await asyncio.to_thread(self.codec.encode, chunk_ndarray)
+        return chunk_spec.prototype.nd_buffer.from_ndarray_like(out)  # type: ignore[arg-type]
+
+
+@dataclass(kw_only=True, frozen=True)
+class NumcodecsArrayBytesCodec(NumcodecsWrapper, ArrayBytesCodec):
+    async def _decode_single(self, chunk_data: Buffer, chunk_spec: ArraySpec) -> NDBuffer:
+        chunk_bytes = chunk_data.to_bytes()
+        out = await asyncio.to_thread(self.codec.decode, chunk_bytes)
+        return chunk_spec.prototype.nd_buffer.from_ndarray_like(out.reshape(chunk_spec.shape))
+
+    async def _encode_single(self, chunk_data: NDBuffer, chunk_spec: ArraySpec) -> Buffer:
+        chunk_ndarray = chunk_data.as_ndarray_like()
+        out = await asyncio.to_thread(self.codec.encode, chunk_ndarray)
+        return chunk_spec.prototype.buffer.from_bytes(out)

src/zarr/codecs/blosc.py

@@ -20,7 +20,6 @@
 from typing_extensions import ReadOnly
 
 from zarr.abc.codec import BytesBytesCodec, CodecJSON
-from zarr.core.buffer.cpu import as_numpy_array_wrapper
 from zarr.core.common import (
     JSON,
     NamedRequiredConfig,
@@ -266,6 +265,7 @@ async def _decode_single(
         chunk_bytes: Buffer,
         chunk_spec: ArraySpec,
     ) -> Buffer:
+        from zarr.core.buffer.cpu import as_numpy_array_wrapper
         return await asyncio.to_thread(
             as_numpy_array_wrapper, self._blosc_codec.decode, chunk_bytes, chunk_spec.prototype
         )

src/zarr/core/_info.py

@@ -4,10 +4,8 @@
 import textwrap
 from typing import TYPE_CHECKING, Literal
 
-from zarr.abc.codec import Numcodec
-
 if TYPE_CHECKING:
-    from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BytesBytesCodec
+    from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BytesBytesCodec, Numcodec
     from zarr.core.common import ZarrFormat
     from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar, ZDType