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

Author: d-v-b

src/zarr/abc/numcodec.py

@@ -0,0 +1,59 @@
+from typing import Self, TypeGuard
+
+from typing_extensions import Protocol
+
+from zarr.abc.codec import CodecJSON_V2
+from zarr.core.buffer import Buffer, NDBuffer
+
+
+class Numcodec(Protocol):
+    """
+    A protocol that models the ``numcodecs.abc.Codec`` interface.
+    """
+
+    codec_id: 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_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

@@ -47,7 +47,7 @@
     from collections.abc import Iterable
 
     from zarr.abc.codec import Codec
-    from zarr.codecs._v2 import Numcodec
+    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

src/zarr/api/synchronous.py

@@ -18,8 +18,8 @@
     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.codecs._v2 import Numcodec
     from zarr.core.array import (
         CompressorsLike,
         FiltersLike,

src/zarr/codecs/_numcodecs.py

@@ -2,73 +2,20 @@
 
 import asyncio
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, ClassVar, Self, TypeGuard
+from typing import TYPE_CHECKING
 
 import numpy as np
 from numcodecs.compat import ensure_bytes, ensure_ndarray_like
-from typing_extensions import Protocol
 
-from zarr.abc.codec import ArrayBytesCodec, CodecJSON_V2
+from zarr.abc.codec import ArrayBytesCodec
 from zarr.registry import get_ndbuffer_class
 
 if TYPE_CHECKING:
+    from zarr.abc.numcodec import Numcodec
     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[Numcodec, ...] | None

src/zarr/codecs/_v2.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

src/zarr/core/_info.py

@@ -19,15 +19,14 @@
 )
 from warnings import warn
 
-import numcodecs
-import numcodecs.abc
 import numpy as np
 from typing_extensions import deprecated
 
 import zarr
 from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BytesBytesCodec, Codec
+from zarr.abc.numcodec import Numcodec, _is_numcodec
 from zarr.abc.store import Store, set_or_delete
-from zarr.codecs._v2 import Numcodec, V2Codec
+from zarr.codecs._v2 import V2Codec
 from zarr.codecs.bytes import BytesCodec
 from zarr.codecs.vlen_utf8 import VLenBytesCodec, VLenUTF8Codec
 from zarr.codecs.zstd import ZstdCodec
@@ -1033,7 +1032,7 @@ def size(self) -> int:
         return np.prod(self.metadata.shape).item()
 
     @property
-    def filters(self) -> tuple[numcodecs.abc.Codec, ...] | tuple[ArrayArrayCodec, ...]:
+    def filters(self) -> tuple[Numcodec, ...] | tuple[ArrayArrayCodec, ...]:
         """
         Filters that are applied to each chunk of the array, in order, before serializing that
         chunk to bytes.
@@ -1062,7 +1061,7 @@ def serializer(self) -> ArrayBytesCodec | None:
 
     @property
     @deprecated("Use AsyncArray.compressors instead.")
-    def compressor(self) -> numcodecs.abc.Codec | None:
+    def compressor(self) -> Numcodec | None:
         """
         Compressor that is applied to each chunk of the array.
 
@@ -1075,7 +1074,7 @@ def compressor(self) -> numcodecs.abc.Codec | None:
         raise TypeError("`compressor` is not available for Zarr format 3 arrays.")
 
     @property
-    def compressors(self) -> tuple[numcodecs.abc.Codec, ...] | tuple[BytesBytesCodec, ...]:
+    def compressors(self) -> tuple[Numcodec, ...] | tuple[BytesBytesCodec, ...]:
         """
         Compressors that are applied to each chunk of the array. Compressors are applied in order, and after any
         filters are applied (if any are specified) and the data is serialized into bytes.
@@ -2227,7 +2226,7 @@ def fill_value(self) -> Any:
         return self.metadata.fill_value
 
     @property
-    def filters(self) -> tuple[numcodecs.abc.Codec, ...] | tuple[ArrayArrayCodec, ...]:
+    def filters(self) -> tuple[Numcodec, ...] | tuple[ArrayArrayCodec, ...]:
         """
         Filters that are applied to each chunk of the array, in order, before serializing that
         chunk to bytes.
@@ -2243,7 +2242,7 @@ def serializer(self) -> None | ArrayBytesCodec:
 
     @property
     @deprecated("Use Array.compressors instead.")
-    def compressor(self) -> numcodecs.abc.Codec | None:
+    def compressor(self) -> Numcodec | None:
         """
         Compressor that is applied to each chunk of the array.
 
@@ -2254,7 +2253,7 @@ def compressor(self) -> numcodecs.abc.Codec | None:
         return self._async_array.compressor
 
     @property
-    def compressors(self) -> tuple[numcodecs.abc.Codec, ...] | tuple[BytesBytesCodec, ...]:
+    def compressors(self) -> tuple[Numcodec, ...] | tuple[BytesBytesCodec, ...]:
         """
         Compressors that are applied to each chunk of the array. Compressors are applied in order, and after any
         filters are applied (if any are specified) and the data is serialized into bytes.
@@ -3900,15 +3899,13 @@ def _build_parents(
 FiltersLike: TypeAlias = (
     Iterable[dict[str, JSON] | ArrayArrayCodec | Numcodec]
     | ArrayArrayCodec
-    | Iterable[numcodecs.abc.Codec]
-    | numcodecs.abc.Codec
+    | Iterable[Numcodec]
+    | Numcodec
     | Literal["auto"]
     | None
 )
 # Union of acceptable types for users to pass in for both v2 and v3 compressors
-CompressorLike: TypeAlias = (
-    dict[str, JSON] | BytesBytesCodec | numcodecs.abc.Codec | Literal["auto"] | None
-)
+CompressorLike: TypeAlias = dict[str, JSON] | BytesBytesCodec | Numcodec | Literal["auto"] | None
 
 CompressorsLike: TypeAlias = (
     Iterable[dict[str, JSON] | BytesBytesCodec | Numcodec]
@@ -4775,7 +4772,7 @@ def default_serializer_v3(dtype: ZDType[Any, Any]) -> ArrayBytesCodec:
     return serializer
 
 
-def default_filters_v2(dtype: ZDType[Any, Any]) -> tuple[numcodecs.abc.Codec] | None:
+def default_filters_v2(dtype: ZDType[Any, Any]) -> tuple[Numcodec] | None:
     """
     Given a data type, return the default filters for that data type.
 
@@ -4797,28 +4794,28 @@ def default_filters_v2(dtype: ZDType[Any, Any]) -> tuple[numcodecs.abc.Codec] |
     return None
 
 
-def default_compressor_v2(dtype: ZDType[Any, Any]) -> numcodecs.abc.Codec:
+def default_compressor_v2(dtype: ZDType[Any, Any]) -> Numcodec:
     """
     Given a data type, return the default compressors for that data type.
 
     This is just the numcodecs ``Zstd`` codec.
     """
     from numcodecs import Zstd
 
-    return Zstd(level=0, checksum=False)
+    return Zstd(level=0, checksum=False)  # type: ignore[no-any-return]
 
 
 def _parse_chunk_encoding_v2(
     *,
     compressor: CompressorsLike,
     filters: FiltersLike,
     dtype: ZDType[TBaseDType, TBaseScalar],
-) -> tuple[tuple[numcodecs.abc.Codec, ...] | None, numcodecs.abc.Codec | None]:
+) -> tuple[tuple[Numcodec, ...] | None, Numcodec | None]:
     """
     Generate chunk encoding classes for Zarr format 2 arrays with optional defaults.
     """
-    _filters: tuple[numcodecs.abc.Codec, ...] | None
-    _compressor: numcodecs.abc.Codec | None
+    _filters: tuple[Numcodec, ...] | None
+    _compressor: Numcodec | None
 
     if compressor is None or compressor == ():
         _compressor = None
@@ -4839,7 +4836,7 @@ def _parse_chunk_encoding_v2(
     else:
         if isinstance(filters, Iterable):
             for idx, f in enumerate(filters):
-                if not isinstance(f, numcodecs.abc.Codec):
+                if not _is_numcodec(f):
                     msg = (
                         "For Zarr format 2 arrays, all elements of `filters` must be numcodecs codecs. "
                         f"Element at index {idx} has type {type(f)}, which is not a numcodecs codec."
@@ -4852,12 +4849,12 @@ def _parse_chunk_encoding_v2(
             if _compressor is None:
                 object_codec_id = None
             else:
-                object_codec_id = get_object_codec_id((_compressor.get_config(),))
+                object_codec_id = get_object_codec_id((_compressor.get_config(),))  # type: ignore[arg-type]
         else:
             object_codec_id = get_object_codec_id(
                 (
-                    *[f.get_config() for f in _filters],
-                    _compressor.get_config() if _compressor is not None else None,
+                    *[f.get_config() for f in _filters],  # type: ignore[arg-type]
+                    _compressor.get_config() if _compressor is not None else None,  # type: ignore[arg-type]
                 )
             )
         if object_codec_id is None:
@@ -4944,7 +4941,7 @@ def _parse_deprecated_compressor(
             # "no compression"
             compressors = ()
         else:
-            compressors = (compressor,)  # type: ignore[assignment]
+            compressors = (compressor,)
     elif zarr_format == 2 and compressor == compressors == "auto":
         compressors = ({"id": "blosc"},)
     return compressors