improve method annotations for numcodecs codecs

Author: d-v-b

src/zarr/abc/codec.py

@@ -8,12 +8,11 @@
     Literal,
     Self,
     TypedDict,
-    TypeGuard,
     TypeVar,
     overload,
 )
 
-from typing_extensions import ReadOnly
+from typing_extensions import ReadOnly, TypeIs
 
 from zarr.abc.metadata import Metadata
 from zarr.core.buffer import Buffer, NDBuffer
@@ -46,16 +45,14 @@
 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]):
+class CodecJSON_V2(TypedDict):
     """The JSON representation of a codec for Zarr V2"""
 
-    id: ReadOnly[TName]
+    id: ReadOnly[str]
 
 
-def _check_codecjson_v2(data: object) -> TypeGuard[CodecJSON_V2[str]]:
+def _check_codecjson_v2(data: object) -> TypeIs[CodecJSON_V2]:
     return isinstance(data, Mapping) and "id" in data and isinstance(data["id"], str)
 
 
@@ -64,7 +61,7 @@ def _check_codecjson_v2(data: object) -> TypeGuard[CodecJSON_V2[str]]:
 
 # The widest type we will *accept* for a codec JSON
 # This covers v2 and v3
-CodecJSON = str | Mapping[str, object]
+CodecJSON = CodecJSON_V2 | CodecJSON_V3
 """The widest type of JSON-like input that could specify a codec."""
 
 
@@ -191,34 +188,28 @@ async def encode(
         return await _batching_helper(self._encode_single, chunks_and_specs)
 
     @overload
-    def to_json(self, zarr_format: Literal[2]) -> CodecJSON_V2[str]: ...
+    def to_json(self, zarr_format: Literal[2]) -> CodecJSON_V2: ...
     @overload
-    def to_json(self, zarr_format: Literal[3]) -> NamedConfig[str, Mapping[str, object]]: ...
+    def to_json(self, zarr_format: Literal[3]) -> CodecJSON_V3: ...
 
-    def to_json(
-        self, zarr_format: ZarrFormat
-    ) -> CodecJSON_V2[str] | NamedConfig[str, Mapping[str, object]]:
+    def to_json(self, zarr_format: ZarrFormat) -> CodecJSON_V2 | CodecJSON_V3:
         raise NotImplementedError
 
     @classmethod
-    def _from_json_v2(cls, data: CodecJSON_V2[str]) -> Self:
+    def _from_json_v2(cls, data: CodecJSON_V2) -> Self:
         return cls(**{k: v for k, v in data.items() if k != "id"})
 
     @classmethod
     def _from_json_v3(cls, data: CodecJSON_V3) -> Self:
         if isinstance(data, str):
             return cls()
-        return cls(**data["configuration"])
+        return cls(**data.get("configuration", {}))
 
     @classmethod
-    def from_json(cls, data: CodecJSON, zarr_format: ZarrFormat) -> Self:
-        if zarr_format == 2:
+    def from_json(cls, data: CodecJSON) -> Self:
+        if _check_codecjson_v2(data):
             return cls._from_json_v2(data)
-        elif zarr_format == 3:
-            return cls._from_json_v3(data)
-        raise ValueError(
-            f"Unsupported Zarr format {zarr_format}. Expected 2 or 3."
-        )  # pragma: no cover
+        return cls._from_json_v3(data)
 
 
 class ArrayArrayCodec(BaseCodec[NDBuffer, NDBuffer]):

src/zarr/abc/numcodec.py

@@ -1,8 +1,9 @@
 from typing import Any, Self, TypeGuard
 
-from typing_extensions import Protocol
+from typing_extensions import Protocol, runtime_checkable
 
 
+@runtime_checkable
 class Numcodec(Protocol):
     """
     A protocol that models the ``numcodecs.abc.Codec`` interface.
@@ -88,14 +89,3 @@ def _is_numcodec_cls(obj: object) -> TypeGuard[type[Numcodec]]:
         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/codecs/_v2.py

@@ -2,8 +2,7 @@
 
 import asyncio
 from dataclasses import dataclass
-from functools import cached_property
-from typing import TYPE_CHECKING, Literal, Self, overload
+from typing import TYPE_CHECKING, Literal, Self, cast, overload
 
 import numpy as np
 from numcodecs.compat import ensure_bytes, ensure_ndarray_like
@@ -15,11 +14,9 @@
     CodecJSON,
     CodecJSON_V2,
 )
-from zarr.registry import get_ndbuffer_class
+from zarr.registry import _get_codec_v2, _get_codec_v3, get_ndbuffer_class
 
 if TYPE_CHECKING:
-    from collections.abc import Mapping
-
     from zarr.abc.numcodec import Numcodec
     from zarr.core.array_spec import ArraySpec
     from zarr.core.buffer import Buffer, NDBuffer
@@ -115,21 +112,16 @@ def compute_encoded_size(self, _input_byte_length: int, _chunk_spec: ArraySpec)
 
 @dataclass(frozen=True, kw_only=True)
 class NumcodecsWrapper:
-    codec_cls: type[Numcodec]
-    config: Mapping[str, object]
-
-    @cached_property
-    def codec(self) -> Numcodec:
-        return self.codec_cls(**self.config)
+    codec: Numcodec
 
     @overload
-    def to_json(self, zarr_format: Literal[2]) -> CodecJSON_V2[str]: ...
+    def to_json(self, zarr_format: Literal[2]) -> CodecJSON_V2: ...
     @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]:
+    def to_json(self, zarr_format: ZarrFormat) -> CodecJSON_V2 | NamedConfig[str, BaseConfig]:
         if zarr_format == 2:
-            return self.config
+            return cast(CodecJSON_V2, 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"}
@@ -138,11 +130,13 @@ def to_json(self, zarr_format: ZarrFormat) -> CodecJSON_V2[str] | NamedConfig[st
 
     @classmethod
     def _from_json_v2(cls, data: CodecJSON) -> Self:
-        return cls(config=data)
+        codec = _get_codec_v2(data)
+        return cls(codec=codec)
 
     @classmethod
     def _from_json_v3(cls, data: CodecJSON) -> Self:
-        return cls(config=data.get("configuration", {}))
+        codec = _get_codec_v3(data)
+        return cls(codec=codec)
 
     def compute_encoded_size(self, input_byte_length: int, chunk_spec: ArraySpec) -> int:
         raise NotImplementedError
@@ -185,19 +179,19 @@ def to_array_array(self) -> NumcodecsArrayArrayCodec:
         """
         Use the ``_codec`` attribute to create a NumcodecsArrayArrayCodec.
         """
-        return NumcodecsArrayArrayCodec(cls=self.codec_cls, config=self.config)
+        return NumcodecsArrayArrayCodec(codec=self.codec)
 
     def to_bytes_bytes(self) -> NumcodecsBytesBytesCodec:
         """
         Use the ``_codec`` attribute to create a NumcodecsBytesBytesCodec.
         """
-        return NumcodecsBytesBytesCodec(cls=self.codec_cls, config=self.config)
+        return NumcodecsBytesBytesCodec(codec=self.codec)
 
     def to_array_bytes(self) -> NumcodecsArrayBytesCodec:
         """
         Use the ``_codec`` attribute to create a NumcodecsArrayBytesCodec.
         """
-        return NumcodecsArrayBytesCodec(codec_cls=self.codec_cls, config=self.config)
+        return NumcodecsArrayBytesCodec(codec=self.codec)
 
 
 class NumcodecsBytesBytesCodec(NumcodecsWrapper, BytesBytesCodec):
@@ -226,12 +220,12 @@ 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]
+        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) -> 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]
+        return chunk_spec.prototype.nd_buffer.from_ndarray_like(out)
 
 
 @dataclass(kw_only=True, frozen=True)

src/zarr/codecs/blosc.py

@@ -75,17 +75,17 @@ def check_json_v2(data: CodecJSON) -> TypeGuard[BloscJSON_V2]:
     return (
         isinstance(data, Mapping)
         and set(data.keys()) == {"id", "clevel", "cname", "shuffle", "blocksize"}
-        and data["id"] == "blosc"
+        and data["id"] == "blosc"  # type: ignore[typeddict-item]
     )
 
 
 def check_json_v3(data: CodecJSON) -> TypeGuard[BloscJSON_V3]:
     return (
         isinstance(data, Mapping)
         and set(data.keys()) == {"name", "configuration"}
-        and data["name"] == "blosc"
-        and isinstance(data["configuration"], Mapping)
-        and set(data["configuration"].keys())
+        and data["name"] == "blosc"  # type: ignore[typeddict-item]
+        and isinstance(data["configuration"], Mapping)  # type: ignore[typeddict-item]
+        and set(data["configuration"].keys())  # type: ignore[typeddict-item]
         == {"cname", "clevel", "shuffle", "blocksize", "typesize"}
     )
 

src/zarr/codecs/bytes.py

@@ -4,7 +4,7 @@
 from collections.abc import Mapping
 from dataclasses import dataclass, replace
 from enum import Enum
-from typing import TYPE_CHECKING, Final, Literal, NotRequired, TypedDict, TypeGuard, overload
+from typing import TYPE_CHECKING, Final, Literal, NotRequired, TypedDict, TypeGuard, cast, overload
 
 import numpy as np
 from typing_extensions import ReadOnly
@@ -61,7 +61,7 @@ def check_json_v2(data: CodecJSON) -> TypeGuard[BytesJSON_V2]:
     return (
         isinstance(data, Mapping)
         and set(data.keys()) in ({"id", "endian"}, {"id"})
-        and data["id"] == "bytes"
+        and data["id"] == "bytes"  # type: ignore[typeddict-item]
     )
 
 
@@ -70,10 +70,10 @@ def check_json_v3(data: CodecJSON) -> TypeGuard[BytesJSON_V3]:
         (
             isinstance(data, Mapping)
             and set(data.keys()) in ({"name"}, {"name", "configuration"})
-            and data["name"] == "bytes"
+            and data["name"] == "bytes"  # type: ignore[typeddict-item]
         )
         and isinstance(data.get("configuration", {}), Mapping)
-        and set(data.get("configuration", {}).keys()) in ({"endian"}, set())
+        and set(data.get("configuration", {}).keys()) in ({"endian"}, set())  # type: ignore[attr-defined]
     )
 
 
@@ -90,10 +90,10 @@ def __init__(self, *, endian: EndiannessStr | str | None = default_system_endian
 
     @classmethod
     def from_dict(cls, data: dict[str, JSON]) -> Self:
-        return cls.from_json(data, zarr_format=3)
+        return cls.from_json(data)  # type: ignore[arg-type]
 
     def to_dict(self) -> dict[str, JSON]:
-        return self.to_json(zarr_format=3)
+        return cast(dict[str, JSON], self.to_json(zarr_format=3))
 
     @classmethod
     def _from_json_v2(cls, data: CodecJSON) -> Self:
@@ -108,7 +108,7 @@ def _from_json_v3(cls, data: CodecJSON) -> Self:
             if data in ("bytes", {"name": "bytes"}, {"name": "bytes", "configuration": {}}):
                 return cls()
             else:
-                return cls(endian=data["configuration"].get("endian", None))
+                return cls(endian=data["configuration"].get("endian", None))  # type: ignore[union-attr, index]
         raise ValueError(f"Invalid JSON: {data}")
 
     @overload
@@ -182,11 +182,11 @@ async def _encode_single(
         if (
             chunk_array.dtype.itemsize > 1
             and self.endian is not None
-            and self.endian != chunk_array.byteorder
+            and self.endian != chunk_array.byteorder.value
         ):
             # type-ignore is a numpy bug
             # see https://github.com/numpy/numpy/issues/26473
-            new_dtype = chunk_array.dtype.newbyteorder(self.endian)  # type: ignore[arg-type]
+            new_dtype = chunk_array.dtype.newbyteorder(self.endian)
             chunk_array = chunk_array.astype(new_dtype)
 
         nd_array = chunk_array.as_ndarray_like()

src/zarr/codecs/crc32c_.py

@@ -7,8 +7,9 @@
 import numpy as np
 import typing_extensions
 from crc32c import crc32c
+from typing_extensions import ReadOnly
 
-from zarr.abc.codec import BytesBytesCodec, CodecJSON, CodecJSON_V2
+from zarr.abc.codec import BytesBytesCodec, CodecJSON
 from zarr.core.common import JSON, NamedConfig, ZarrFormat, parse_named_configuration
 from zarr.errors import CodecValidationError
 
@@ -19,13 +20,14 @@
     from zarr.core.buffer import Buffer
 
 
-class Crc32Config(TypedDict): ...
+class Crc32cConfig(TypedDict): ...
 
 
-class Crc32cJSON_V2(CodecJSON_V2[Literal["crc32c"]]): ...
+class Crc32cJSON_V2(TypedDict):
+    id: ReadOnly[Literal["crc32c"]]
 
 
-class Crc32cJSON_V3(NamedConfig[Literal["crc32c"], Crc32Config]): ...
+class Crc32cJSON_V3(NamedConfig[Literal["crc32c"], Crc32cConfig]): ...
 
 
 def check_json_v2(data: CodecJSON) -> TypeGuard[Crc32cJSON_V2]:
@@ -47,7 +49,6 @@ class Crc32cCodec(BytesBytesCodec):
 
     @classmethod
     def from_dict(cls, data: dict[str, JSON]) -> Self:
-        return cls.from_json(data, zarr_format=3)
         parse_named_configuration(data, "crc32c", require_configuration=False)
         return cls()
 
@@ -72,7 +73,6 @@ def _from_json_v3(cls, data: CodecJSON) -> Self:
         raise CodecValidationError(msg)
 
     def to_dict(self) -> dict[str, JSON]:
-        return self.to_json(zarr_format=3)
         return {"name": "crc32c"}
 
     @overload

src/zarr/codecs/gzip.py

@@ -3,7 +3,7 @@
 import asyncio
 from collections.abc import Mapping
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Literal, TypedDict, TypeGuard, overload
+from typing import TYPE_CHECKING, Literal, TypedDict, TypeGuard, cast, overload
 
 from numcodecs.gzip import GZip
 from typing_extensions import ReadOnly
@@ -64,10 +64,10 @@ def __init__(self, *, level: int = 5) -> None:
 
     @classmethod
     def from_dict(cls, data: dict[str, JSON]) -> Self:
-        return cls.from_json(data, zarr_format=3)
+        return cls.from_json(data)  # type: ignore[arg-type]
 
     def to_dict(self) -> dict[str, JSON]:
-        return self.to_json(zarr_format=3)
+        return cast(dict[str, JSON], self.to_json(zarr_format=3))
 
     @overload
     def to_json(self, zarr_format: Literal[2]) -> GZipJSON_V2: ...
@@ -88,19 +88,19 @@ def _check_json_v2(cls, data: CodecJSON) -> TypeGuard[GZipJSON_V2]:
         return (
             isinstance(data, Mapping)
             and set(data.keys()) == {"id", "level"}
-            and data["id"] == "gzip"
-            and isinstance(data["level"], int)
+            and data["id"] == "gzip"  # type: ignore[typeddict-item]
+            and isinstance(data["level"], int)  # type: ignore[typeddict-item]
         )
 
     @classmethod
     def _check_json_v3(cls, data: CodecJSON) -> TypeGuard[GZipJSON_V3]:
         return (
             isinstance(data, Mapping)
             and set(data.keys()) == {"name", "configuration"}
-            and data["name"] == "gzip"
-            and isinstance(data["configuration"], dict)
-            and "level" in data["configuration"]
-            and isinstance(data["configuration"]["level"], int)
+            and data["name"] == "gzip"  # type: ignore[typeddict-item]
+            and isinstance(data["configuration"], Mapping)  # type: ignore[typeddict-item]
+            and "level" in data["configuration"]  # type: ignore[typeddict-item]
+            and isinstance(data["configuration"]["level"], int)  # type: ignore[typeddict-item]
         )
 
     @classmethod