add numcodec protocol (#3318)

* add numcodec protocol

* add tests for numcodecs compatibility

* changelog

* ignore unknown key

* remove re-implementation of get_codec

* avoid circular imports by importing lower-level routines exactly where needed

* push numcodec prototol into abcs; remove all numcodecs.abc.Codec type annotations

* add tests for codecjson typeguard

* avoid using zarr's buffer / ndbuffer for numcodec encode / decode

* use Any to model input / output types of numcodec protocol

* add numcodec protocol

* add tests for numcodecs compatibility

* changelog

* ignore unknown key

* remove re-implementation of get_codec

* avoid circular imports by importing lower-level routines exactly where needed

* push numcodec prototol into abcs; remove all numcodecs.abc.Codec type annotations

* add tests for codecjson typeguard

* avoid using zarr's buffer / ndbuffer for numcodec encode / decode

* use Any to model input / output types of numcodec protocol

* Update src/zarr/abc/numcodec.py

Co-authored-by: Max Jones <[email protected]>

* Update src/zarr/abc/numcodec.py

Co-authored-by: Max Jones <[email protected]>

* Update src/zarr/abc/numcodec.py

Co-authored-by: Max Jones <[email protected]>

* Update src/zarr/abc/numcodec.py

Co-authored-by: Max Jones <[email protected]>

* Update src/zarr/abc/numcodec.py

Co-authored-by: Max Jones <[email protected]>

* fix docstrings

* revert changes to store imports

* remove whitespace

* fix docstring

---------

Co-authored-by: Max Jones <[email protected]>

Author: d-v-b

changes/3318.misc.rst

@@ -0,0 +1,2 @@
+Define a ``Protocol`` to model the ``numcodecs.abc.Codec`` interface. This is groundwork toward
+making ``numcodecs`` an optional dependency for ``zarr-python``.

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/abc/numcodec.py

@@ -0,0 +1,101 @@
+from typing import Any, Self, TypeGuard
+
+from typing_extensions import Protocol
+
+
+class Numcodec(Protocol):
+    """
+    A protocol that models the ``numcodecs.abc.Codec`` interface.
+
+    This protocol should be considered experimental. Expect the type annotations for ``buf`` and
+    ``out`` to narrow in the future.
+    """
+
+    codec_id: str
+
+    def encode(self, buf: Any) -> Any:
+        """Encode data from ``buf``.
+
+        Parameters
+        ----------
+        buf : Any
+            Data to be encoded.
+
+        Returns
+        -------
+        enc: Any
+            Encoded data.
+        """
+        ...
+
+    def decode(self, buf: Any, out: Any | None = None) -> Any:
+        """
+        Decode data in ``buf``.
+
+        Parameters
+        ----------
+        buf : Any
+            Encoded data.
+        out : Any
+            Writeable buffer to store decoded data. If provided, this buffer must
+            be exactly the right size to store the decoded data.
+
+        Returns
+        -------
+        dec : Any
+            Decoded data.
+        """
+        ...
+
+    def get_config(self) -> Any:
+        """
+        Return a JSON-serializable configuration dictionary for this
+        codec. Must include an ``'id'`` field with the codec identifier.
+        """
+        ...
+
+    @classmethod
+    def from_config(cls, config: Any) -> Self:
+        """
+        Instantiate a codec from a configuration dictionary.
+
+        Parameters
+        ----------
+        config : Any
+            A configuration dictionary for this codec.
+        """
+        ...
+
+
+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)
+    )
+
+
+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))

src/zarr/api/asynchronous.py

@@ -52,9 +52,8 @@
 if TYPE_CHECKING:
     from collections.abc import Iterable
 
-    import numcodecs.abc
-
     from zarr.abc.codec import Codec
+    from zarr.abc.numcodec import Numcodec
     from zarr.core.buffer import NDArrayLikeOrScalar
     from zarr.core.chunk_key_encodings import ChunkKeyEncoding
     from zarr.storage import StoreLike
@@ -877,7 +876,7 @@ async def create(
     overwrite: bool = False,
     path: PathLike | None = None,
     chunk_store: StoreLike | None = None,
-    filters: Iterable[dict[str, JSON] | numcodecs.abc.Codec] | None = None,
+    filters: Iterable[dict[str, JSON] | Numcodec] | None = None,
     cache_metadata: bool | None = None,
     cache_attrs: bool | None = None,
     read_only: bool | None = None,

src/zarr/api/synchronous.py

@@ -15,11 +15,11 @@
 if TYPE_CHECKING:
     from collections.abc import Iterable
 
-    import numcodecs.abc
     import numpy as np
     import numpy.typing as npt
 
     from zarr.abc.codec import Codec
+    from zarr.abc.numcodec import Numcodec
     from zarr.api.asynchronous import ArrayLike, PathLike
     from zarr.core.array import (
         CompressorsLike,
@@ -610,7 +610,7 @@ def create(
     overwrite: bool = False,
     path: PathLike | None = None,
     chunk_store: StoreLike | None = None,
-    filters: Iterable[dict[str, JSON] | numcodecs.abc.Codec] | None = None,
+    filters: Iterable[dict[str, JSON] | Numcodec] | None = None,
     cache_metadata: bool | None = None,
     cache_attrs: bool | None = None,
     read_only: bool | None = None,

src/zarr/codecs/_v2.py

@@ -4,24 +4,22 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
-import numcodecs
 import numpy as np
 from numcodecs.compat import ensure_bytes, ensure_ndarray_like
 
 from zarr.abc.codec import ArrayBytesCodec
 from zarr.registry import get_ndbuffer_class
 
 if TYPE_CHECKING:
-    import numcodecs.abc
-
+    from zarr.abc.numcodec import Numcodec
     from zarr.core.array_spec import ArraySpec
     from zarr.core.buffer import Buffer, NDBuffer
 
 
 @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
 
@@ -86,7 +84,6 @@ async def _encode_single(
         if self.filters:
             for f in self.filters:
                 chunk = await asyncio.to_thread(f.encode, chunk)
-
         # check object encoding
         if ensure_ndarray_like(chunk).dtype == object:
             raise RuntimeError("cannot write object array without object codec")
@@ -96,7 +93,6 @@ async def _encode_single(
             cdata = await asyncio.to_thread(self.compressor.encode, chunk)
         else:
             cdata = chunk
-
         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.abc.numcodec 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