add numcodecs protocol

Author: d-v-b

src/zarr/abc/codec.py

@@ -1,11 +1,14 @@
 from __future__ import annotations
 
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Generic, TypeVar
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Generic, TypeGuard, TypeVar
+
+from typing_extensions import ReadOnly, TypedDict
 
 from zarr.abc.metadata import Metadata
 from zarr.core.buffer import Buffer, NDBuffer
-from zarr.core.common import ChunkCoords, concurrent_map
+from zarr.core.common import ChunkCoords, NamedConfig, concurrent_map
 from zarr.core.config import config
 
 if TYPE_CHECKING:
@@ -34,6 +37,27 @@
 CodecInput = TypeVar("CodecInput", bound=NDBuffer | Buffer)
 CodecOutput = TypeVar("CodecOutput", bound=NDBuffer | Buffer)
 
+TName = TypeVar("TName", bound=str, covariant=True)
+
+
+class CodecJSON_V2(TypedDict, Generic[TName]):
+    """The JSON representation of a codec for Zarr V2"""
+
+    id: ReadOnly[TName]
+
+
+def _check_codecjson_v2(data: object) -> TypeGuard[CodecJSON_V2[str]]:
+    return isinstance(data, Mapping) and "id" in data and isinstance(data["id"], str)
+
+
+CodecJSON_V3 = str | NamedConfig[str, Mapping[str, object]]
+"""The JSON representation of a codec for Zarr V3."""
+
+# The widest type we will *accept* for a codec JSON
+# This covers v2 and v3
+CodecJSON = str | Mapping[str, object]
+"""The widest type of JSON-like input that could specify a codec."""
+
 
 class BaseCodec(Metadata, Generic[CodecInput, CodecOutput]):
     """Generic base class for codecs.

src/zarr/codecs/_numcodecs.py

@@ -0,0 +1,37 @@
+import numcodecs.registry as numcodecs_registry
+
+from zarr.abc.codec import CodecJSON_V2
+from zarr.codecs._v2 import Numcodec
+
+
+def get_numcodec(data: CodecJSON_V2[str]) -> Numcodec:
+    """
+    Resolve a numcodec codec from the numcodecs registry.
+
+    This requires the Numcodecs package to be installed.
+
+    Parameters
+    ----------
+    data : CodecJSON_V2
+        The JSON metadata for the codec.
+
+    Returns
+    -------
+    codec : Numcodec
+
+    Examples
+    --------
+
+    >>> codec = get_numcodec({'id': 'zlib', 'level': 1})
+    >>> codec
+    Zlib(level=1)
+    """
+
+    codec_id = data["id"]
+    cls = numcodecs_registry.codec_registry.get(codec_id)
+    if cls is None and data in numcodecs_registry.entries:
+        cls = numcodecs_registry.entries[data].load()
+        numcodecs_registry.register_codec(cls, codec_id=data)
+    if cls is not None:
+        return cls.from_config({k: v for k, v in data.items() if k != "id"})  # type: ignore[no-any-return]
+    raise KeyError(data)

src/zarr/codecs/_v2.py

@@ -2,26 +2,77 @@
 
 import asyncio
 from dataclasses import dataclass
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, ClassVar, Self, TypeGuard
 
-import numcodecs
 import numpy as np
 from numcodecs.compat import ensure_bytes, ensure_ndarray_like
+from typing_extensions import Protocol
 
-from zarr.abc.codec import ArrayBytesCodec
+from zarr.abc.codec import ArrayBytesCodec, CodecJSON_V2
 from zarr.registry import get_ndbuffer_class
 
 if TYPE_CHECKING:
-    import numcodecs.abc
-
     from zarr.core.array_spec import ArraySpec
     from zarr.core.buffer import Buffer, NDBuffer
 
 
+class Numcodec(Protocol):
+    """
+    A protocol that models the ``numcodecs.abc.Codec`` interface.
+    """
+
+    codec_id: ClassVar[str]
+
+    def encode(self, buf: Buffer | NDBuffer) -> Buffer | NDBuffer: ...
+
+    def decode(
+        self, buf: Buffer | NDBuffer, out: Buffer | NDBuffer | None = None
+    ) -> Buffer | NDBuffer: ...
+
+    def get_config(self) -> CodecJSON_V2[str]: ...
+
+    @classmethod
+    def from_config(cls, config: CodecJSON_V2[str]) -> Self: ...
+
+
+def _is_numcodec(obj: object) -> TypeGuard[Numcodec]:
+    """
+    Check if the given object implements the Numcodec protocol.
+
+    The @runtime_checkable decorator does not allow issubclass checks for protocols with non-method
+    members (i.e., attributes), so we use this function to manually check for the presence of the
+    required attributes and methods on a given object.
+    """
+    return _is_numcodec_cls(type(obj))
+
+
+def _is_numcodec_cls(obj: object) -> TypeGuard[type[Numcodec]]:
+    """
+    Check if the given object is a class implements the Numcodec protocol.
+
+    The @runtime_checkable decorator does not allow issubclass checks for protocols with non-method
+    members (i.e., attributes), so we use this function to manually check for the presence of the
+    required attributes and methods on a given object.
+    """
+    return (
+        isinstance(obj, type)
+        and hasattr(obj, "codec_id")
+        and isinstance(obj.codec_id, str)
+        and hasattr(obj, "encode")
+        and callable(obj.encode)
+        and hasattr(obj, "decode")
+        and callable(obj.decode)
+        and hasattr(obj, "get_config")
+        and callable(obj.get_config)
+        and hasattr(obj, "from_config")
+        and callable(obj.from_config)
+    )
+
+
 @dataclass(frozen=True)
 class V2Codec(ArrayBytesCodec):
-    filters: tuple[numcodecs.abc.Codec, ...] | None
-    compressor: numcodecs.abc.Codec | None
+    filters: tuple[Numcodec, ...] | None
+    compressor: Numcodec | None
 
     is_fixed_size = False
 
@@ -33,9 +84,9 @@ async def _decode_single(
         cdata = chunk_bytes.as_array_like()
         # decompress
         if self.compressor:
-            chunk = await asyncio.to_thread(self.compressor.decode, cdata)
+            chunk = await asyncio.to_thread(self.compressor.decode, cdata)  # type: ignore[arg-type]
         else:
-            chunk = cdata
+            chunk = cdata  # type: ignore[assignment]
 
         # apply filters
         if self.filters:
@@ -56,7 +107,7 @@ async def _decode_single(
                 # is an object array. In this case, we need to convert the object
                 # array to the correct dtype.
 
-                chunk = np.array(chunk).astype(chunk_spec.dtype.to_native_dtype())
+                chunk = np.array(chunk).astype(chunk_spec.dtype.to_native_dtype())  # type: ignore[assignment]
 
         elif chunk.dtype != object:
             # If we end up here, someone must have hacked around with the filters.
@@ -85,17 +136,17 @@ async def _encode_single(
         # apply filters
         if self.filters:
             for f in self.filters:
-                chunk = await asyncio.to_thread(f.encode, chunk)
+                chunk = await asyncio.to_thread(f.encode, chunk)  # type: ignore[arg-type]
 
         # check object encoding
         if ensure_ndarray_like(chunk).dtype == object:
             raise RuntimeError("cannot write object array without object codec")
 
         # compress
         if self.compressor:
-            cdata = await asyncio.to_thread(self.compressor.encode, chunk)
+            cdata = await asyncio.to_thread(self.compressor.encode, chunk)  # type: ignore[arg-type]
         else:
-            cdata = chunk
+            cdata = chunk  # type: ignore[assignment]
 
         cdata = ensure_bytes(cdata)
         return chunk_spec.prototype.buffer.from_bytes(cdata)

src/zarr/core/_info.py

@@ -5,9 +5,8 @@
 from typing import TYPE_CHECKING, Literal
 
 if TYPE_CHECKING:
-    import numcodecs.abc
-
     from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BytesBytesCodec
+    from zarr.codecs._v2 import Numcodec
     from zarr.core.common import ZarrFormat
     from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar, ZDType
 
@@ -88,9 +87,9 @@ class ArrayInfo:
     _order: Literal["C", "F"]
     _read_only: bool
     _store_type: str
-    _filters: tuple[numcodecs.abc.Codec, ...] | tuple[ArrayArrayCodec, ...] = ()
+    _filters: tuple[Numcodec, ...] | tuple[ArrayArrayCodec, ...] = ()
     _serializer: ArrayBytesCodec | None = None
-    _compressors: tuple[numcodecs.abc.Codec, ...] | tuple[BytesBytesCodec, ...] = ()
+    _compressors: tuple[Numcodec, ...] | tuple[BytesBytesCodec, ...] = ()
     _count_bytes: int | None = None
     _count_bytes_stored: int | None = None
     _count_chunks_initialized: int | None = None

src/zarr/core/metadata/v2.py

@@ -1,13 +1,14 @@
 from __future__ import annotations
 
 import warnings
-from collections.abc import Iterable, Sequence
+from collections.abc import Iterable, Mapping, Sequence
 from functools import cached_property
 from typing import TYPE_CHECKING, Any, TypeAlias, TypedDict, cast
 
-import numcodecs.abc
-
+from zarr.abc.codec import CodecJSON_V2, _check_codecjson_v2
 from zarr.abc.metadata import Metadata
+from zarr.codecs._numcodecs import get_numcodec
+from zarr.codecs._v2 import Numcodec, _is_numcodec
 from zarr.core.chunk_grids import RegularChunkGrid
 from zarr.core.dtype import get_data_type_from_json
 from zarr.core.dtype.common import OBJECT_CODEC_IDS, DTypeSpec_V2
@@ -30,7 +31,6 @@
 import json
 from dataclasses import dataclass, field, fields, replace
 
-import numcodecs
 import numpy as np
 
 from zarr.core.array_spec import ArrayConfig, ArraySpec
@@ -56,7 +56,7 @@ class ArrayV2MetadataDict(TypedDict):
 
 
 # Union of acceptable types for v2 compressors
-CompressorLikev2: TypeAlias = dict[str, JSON] | numcodecs.abc.Codec | None
+CompressorLike_V2: TypeAlias = CodecJSON_V2[str] | Numcodec
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -66,9 +66,9 @@ class ArrayV2Metadata(Metadata):
     dtype: ZDType[TBaseDType, TBaseScalar]
     fill_value: int | float | str | bytes | None = None
     order: MemoryOrder = "C"
-    filters: tuple[numcodecs.abc.Codec, ...] | None = None
+    filters: tuple[Numcodec, ...] | None = None
     dimension_separator: Literal[".", "/"] = "."
-    compressor: numcodecs.abc.Codec | None
+    compressor: Numcodec | None
     attributes: dict[str, JSON] = field(default_factory=dict)
     zarr_format: Literal[2] = field(init=False, default=2)
 
@@ -81,8 +81,8 @@ def __init__(
         fill_value: Any,
         order: MemoryOrder,
         dimension_separator: Literal[".", "/"] = ".",
-        compressor: CompressorLikev2 = None,
-        filters: Iterable[numcodecs.abc.Codec | dict[str, JSON]] | None = None,
+        compressor: CompressorLike_V2 | None = None,
+        filters: Iterable[CompressorLike_V2] | None = None,
         attributes: dict[str, JSON] | None = None,
     ) -> None:
         """
@@ -197,12 +197,12 @@ def from_dict(cls, data: dict[str, Any]) -> ArrayV2Metadata:
 
     def to_dict(self) -> dict[str, JSON]:
         zarray_dict = super().to_dict()
-        if isinstance(zarray_dict["compressor"], numcodecs.abc.Codec):
+        if _is_numcodec(zarray_dict["compressor"]):
             codec_config = zarray_dict["compressor"].get_config()
             # Hotfix for https://github.com/zarr-developers/zarr-python/issues/2647
-            if codec_config["id"] == "zstd" and not codec_config.get("checksum", False):
-                codec_config.pop("checksum")
-            zarray_dict["compressor"] = codec_config
+            if codec_config.get("id") == "zstd" and not codec_config.get("checksum", False):
+                codec_config.pop("checksum")  # type: ignore[typeddict-item]
+            zarray_dict["compressor"] = codec_config  # type: ignore[assignment]
 
         if zarray_dict["filters"] is not None:
             raw_filters = zarray_dict["filters"]
@@ -212,11 +212,12 @@ def to_dict(self) -> dict[str, JSON]:
                 raise TypeError("Invalid type for filters. Expected a list or tuple.")
             new_filters = []
             for f in raw_filters:
-                if isinstance(f, numcodecs.abc.Codec):
+                if _is_numcodec(f):
                     new_filters.append(f.get_config())
                 else:
                     new_filters.append(f)
-            zarray_dict["filters"] = new_filters
+            # TODO: remove the type ignore when we property type the output
+            zarray_dict["filters"] = new_filters  # type: ignore[assignment]
 
         # serialize the fill value after dtype-specific JSON encoding
         if self.fill_value is not None:
@@ -262,44 +263,44 @@ def parse_zarr_format(data: object) -> Literal[2]:
     raise ValueError(f"Invalid value. Expected 2. Got {data}.")
 
 
-def parse_filters(data: object) -> tuple[numcodecs.abc.Codec, ...] | None:
+def parse_filters(data: object) -> tuple[Numcodec, ...] | None:
     """
     Parse a potential tuple of filters
     """
-    out: list[numcodecs.abc.Codec] = []
+    out: list[Numcodec] = []
 
     if data is None:
         return data
     if isinstance(data, Iterable):
-        for idx, val in enumerate(data):
-            if isinstance(val, numcodecs.abc.Codec):
+        for val in data:
+            if _is_numcodec(val):
                 out.append(val)
-            elif isinstance(val, dict):
-                out.append(numcodecs.get_codec(val))
+            if _check_codecjson_v2(val):
+                out.append(get_numcodec(val))
             else:
-                msg = f"Invalid filter at index {idx}. Expected a numcodecs.abc.Codec or a dict representation of numcodecs.abc.Codec. Got {type(val)} instead."
+                msg = f'Invalid representation of Numcodec. Got {data}, expected a dict with an "id" key or a Numcodec instance.'
                 raise TypeError(msg)
         if len(out) == 0:
             # Per the v2 spec, an empty tuple is not allowed -- use None to express "no filters"
             return None
         else:
             return tuple(out)
     # take a single codec instance and wrap it in a tuple
-    if isinstance(data, numcodecs.abc.Codec):
+    if _is_numcodec(data):
         return (data,)
-    msg = f"Invalid filters. Expected None, an iterable of numcodecs.abc.Codec or dict representations of numcodecs.abc.Codec. Got {type(data)} instead."
+    msg = f"Invalid filters. Expected None, an iterable of Numcodec or dict representations of Numcodec. Got {type(data)} instead."
     raise TypeError(msg)
 
 
-def parse_compressor(data: object) -> numcodecs.abc.Codec | None:
+def parse_compressor(data: object) -> Numcodec | None:
     """
     Parse a potential compressor.
     """
-    if data is None or isinstance(data, numcodecs.abc.Codec):
+    if data is None or _is_numcodec(data):
         return data
-    if isinstance(data, dict):
-        return numcodecs.get_codec(data)
-    msg = f"Invalid compressor. Expected None, a numcodecs.abc.Codec, or a dict representation of a numcodecs.abc.Codec. Got {type(data)} instead."
+    if _check_codecjson_v2(data):
+        return get_numcodec(data)
+    msg = f"Invalid compressor. Expected None, a Numcodec, or a dict representation of a Numcodec. Got {type(data)} instead."
     raise ValueError(msg)
 
 
@@ -313,7 +314,7 @@ def parse_metadata(data: ArrayV2Metadata) -> ArrayV2Metadata:
     return data
 
 
-def get_object_codec_id(maybe_object_codecs: Sequence[JSON]) -> str | None:
+def get_object_codec_id(maybe_object_codecs: Iterable[object]) -> str | None:
     """
     Inspect a sequence of codecs / filters for an "object codec", i.e. a codec
     that can serialize object arrays to contiguous bytes. Zarr python
@@ -324,7 +325,7 @@ def get_object_codec_id(maybe_object_codecs: Sequence[JSON]) -> str | None:
     object_codec_id = None
     for maybe_object_codec in maybe_object_codecs:
         if (
-            isinstance(maybe_object_codec, dict)
+            isinstance(maybe_object_codec, Mapping)
             and maybe_object_codec.get("id") in OBJECT_CODEC_IDS
         ):
             return cast("str", maybe_object_codec["id"])