pull array-array codecs out into their own modules, and design codec JSON properly

Author: d-v-b

src/zarr/codecs/numcodecs/__init__.py

@@ -10,10 +10,8 @@
     LZMA,
     ZFPY,
     Adler32,
-    AsType,
     BitRound,
     Blosc,
-    FixedScaleOffset,
     Fletcher32,
     GZip,
     JenkinsLookup3,
@@ -28,10 +26,12 @@
     _NumcodecsBytesBytesCodec,
     _NumcodecsCodec,
 )
+from zarr.codecs.numcodecs.astype import AsType
 from zarr.codecs.numcodecs.delta import Delta
+from zarr.codecs.numcodecs.fixed_scale_offset import FixedScaleOffset
 
 # This is a fixed dictionary of numcodecs codecs for which we have pre-made Zarr V3 wrappers
-numcodecs_wrappers: Final[dict[str, type[_NumcodecsCodec]]] = {
+NUMCODECS_WRAPPERS: Final[dict[str, type[_NumcodecsCodec]]] = {
     "bz2": BZ2,
     "crc32": CRC32,
     "crc32c": CRC32C,

src/zarr/codecs/numcodecs/_codecs.py

@@ -61,7 +61,7 @@
     _check_codecjson_v2,
     product,
 )
-from zarr.dtype import UInt8, ZDType, parse_dtype
+from zarr.dtype import UInt8, ZDType
 from zarr.errors import ZarrUserWarning
 from zarr.registry import get_numcodec
 
@@ -75,7 +75,6 @@
     from zarr.codecs.zstd import ZstdConfig_V3, ZstdJSON_V2, ZstdJSON_V3
     from zarr.core.array_spec import ArraySpec
     from zarr.core.buffer import Buffer, BufferPrototype, NDBuffer
-    from zarr.core.dtype.common import DTypeSpec_V2, DTypeSpec_V3
 
 
 # TypedDict definitions for V2 and V3 JSON representations
@@ -159,20 +158,6 @@ class BitRoundConfig(TypedDict):
     keepbits: int
 
 
-class FixedScaleOffsetConfig_V2(TypedDict):
-    dtype: NotRequired[DTypeSpec_V2]
-    astype: NotRequired[DTypeSpec_V2]
-    scale: NotRequired[float]
-    offset: NotRequired[float]
-
-
-class FixedScaleOffsetConfig_V3(TypedDict):
-    dtype: NotRequired[DTypeSpec_V3]
-    astype: NotRequired[DTypeSpec_V2]
-    scale: NotRequired[float]
-    offset: NotRequired[float]
-
-
 class QuantizeConfig(TypedDict):
     digits: int
     dtype: NotRequired[str]
@@ -182,11 +167,6 @@ class PackBitsConfig(TypedDict):
     pass  # PackBits has no configuration parameters
 
 
-class AsTypeConfig(TypedDict):
-    encode_dtype: str
-    decode_dtype: str
-
-
 class BitRoundJSON_V2(BitRoundConfig):
     """JSON representation of BitRound codec for Zarr V2."""
 
@@ -197,18 +177,6 @@ class BitRoundJSON_V3(NamedRequiredConfig[Literal["bitround"], BitRoundConfig]):
     """JSON representation of BitRound codec for Zarr V3."""
 
 
-class FixedScaleOffsetJSON_V2(FixedScaleOffsetConfig_V2):
-    """JSON representation of FixedScaleOffset codec for Zarr V2."""
-
-    id: ReadOnly[Literal["fixedscaleoffset"]]
-
-
-class FixedScaleOffsetJSON_V3(
-    NamedRequiredConfig[Literal["fixedscaleoffset"], FixedScaleOffsetConfig_V3]
-):
-    """JSON representation of FixedScaleOffset codec for Zarr V3."""
-
-
 class QuantizeJSON_V2(QuantizeConfig):
     """JSON representation of Quantize codec for Zarr V2."""
 
@@ -229,16 +197,6 @@ class PackBitsJSON_V3(NamedRequiredConfig[Literal["packbits"], PackBitsConfig]):
     """JSON representation of PackBits codec for Zarr V3."""
 
 
-class AsTypeJSON_V2(AsTypeConfig):
-    """JSON representation of AsType codec for Zarr V2."""
-
-    id: ReadOnly[Literal["astype"]]
-
-
-class AsTypeJSON_V3(NamedRequiredConfig[Literal["astype"], AsTypeConfig]):
-    """JSON representation of AsType codec for Zarr V3."""
-
-
 # Checksum codec JSON representations
 class Crc32Config(TypedDict):
     """Configuration parameters for CRC32 codec."""
@@ -601,32 +559,6 @@ def to_json(self, zarr_format: ZarrFormat) -> BitRoundJSON_V2 | BitRoundJSON_V3:
         return super().to_json(zarr_format)  # type: ignore[return-value]
 
 
-class FixedScaleOffset(_NumcodecsArrayArrayCodec):
-    codec_name = "numcodecs.fixedscaleoffset"
-    _codec_id = "fixedscaleoffset"
-    codec_config: FixedScaleOffsetConfig_V2
-
-    @overload
-    def to_json(self, zarr_format: Literal[2]) -> FixedScaleOffsetJSON_V2: ...
-    @overload
-    def to_json(self, zarr_format: Literal[3]) -> FixedScaleOffsetJSON_V3: ...
-    def to_json(self, zarr_format: ZarrFormat) -> FixedScaleOffsetJSON_V2 | FixedScaleOffsetJSON_V3:
-        _warn_unstable_specification(self)
-        return super().to_json(zarr_format)  # type: ignore[return-value]
-
-    def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec:
-        if astype := self.codec_config.get("astype"):
-            dtype = parse_dtype(np.dtype(astype), zarr_format=3)  # type: ignore[arg-type]
-            return replace(chunk_spec, dtype=dtype)
-        return chunk_spec
-
-    def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self:
-        if self.codec_config.get("dtype") is None:
-            dtype = array_spec.dtype.to_native_dtype()
-            return type(self)(**{**self.codec_config, "dtype": str(dtype)})
-        return self
-
-
 class Quantize(_NumcodecsArrayArrayCodec):
     codec_name = "numcodecs.quantize"
     _codec_id = "quantize"
@@ -676,31 +608,6 @@ def validate(self, *, dtype: ZDType[Any, Any], **_kwargs: Any) -> None:
             raise ValueError(f"Packbits filter requires bool dtype. Got {dtype}.")
 
 
-class AsType(_NumcodecsArrayArrayCodec):
-    codec_name = "numcodecs.astype"
-    _codec_id = "astype"
-    codec_config: AsTypeConfig
-
-    @overload
-    def to_json(self, zarr_format: Literal[2]) -> AsTypeJSON_V2: ...
-    @overload
-    def to_json(self, zarr_format: Literal[3]) -> AsTypeJSON_V3: ...
-    def to_json(self, zarr_format: ZarrFormat) -> AsTypeJSON_V2 | AsTypeJSON_V3:
-        _warn_unstable_specification(self)
-        return super().to_json(zarr_format)  # type: ignore[return-value]
-
-    def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec:
-        dtype = parse_dtype(np.dtype(self.codec_config["encode_dtype"]), zarr_format=3)
-        return replace(chunk_spec, dtype=dtype)
-
-    def evolve_from_array_spec(self, array_spec: ArraySpec) -> AsType:
-        if self.codec_config.get("decode_dtype") is None:
-            # TODO: remove these coverage exemptions the correct way, i.e. with tests
-            dtype = array_spec.dtype.to_native_dtype()  # pragma: no cover
-            return AsType(**{**self.codec_config, "decode_dtype": str(dtype)})  # pragma: no cover
-        return self
-
-
 # bytes-to-bytes checksum codecs
 class _NumcodecsChecksumCodec(_NumcodecsBytesBytesCodec):
     def compute_encoded_size(self, input_byte_length: int, chunk_spec: ArraySpec) -> int:

src/zarr/codecs/numcodecs/astype.py

@@ -0,0 +1,126 @@
+from collections.abc import Mapping
+from dataclasses import replace
+from typing import Literal, Self, TypedDict, TypeGuard, overload
+
+from typing_extensions import ReadOnly
+
+from zarr.codecs.numcodecs._codecs import _NumcodecsArrayArrayCodec, _warn_unstable_specification
+from zarr.core.array_spec import ArraySpec
+from zarr.core.common import (
+    CodecJSON,
+    CodecJSON_V2,
+    CodecJSON_V3,
+    NamedRequiredConfig,
+    ZarrFormat,
+    _check_codecjson_v2,
+    _check_codecjson_v3,
+)
+from zarr.core.dtype import parse_dtype
+from zarr.core.dtype.common import (
+    DTypeName_V2,
+    DTypeSpec_V3,
+    check_dtype_name_v2,
+    check_dtype_spec_v3,
+)
+
+
+class AsTypeConfig_V2(TypedDict):
+    encode_dtype: DTypeName_V2
+    decode_dtype: DTypeName_V2
+
+
+class AsTypeConfig_V3(TypedDict):
+    encode_dtype: DTypeSpec_V3
+    decode_dtype: DTypeSpec_V3
+
+
+class AsTypeJSON_V2(AsTypeConfig_V2):
+    """JSON representation of AsType codec for Zarr V2."""
+
+    id: ReadOnly[Literal["astype"]]
+
+
+class AsTypeJSON_V3(NamedRequiredConfig[Literal["astype"], AsTypeConfig_V3]):
+    """JSON representation of AsType codec for Zarr V3."""
+
+
+def check_json_v2(data: object) -> TypeGuard[AsTypeJSON_V2]:
+    """
+    A type guard for the Zarr V2 form of the Astype codec JSON
+    """
+    return (
+        _check_codecjson_v2(data)
+        and data["id"] == "astype"
+        and "encode_dtype" in data
+        and "decode_dtype" in data
+        and check_dtype_name_v2(data["encode_dtype"])  # type: ignore[typeddict-item]
+        and check_dtype_name_v2(data["decode_dtype"])  # type: ignore[typeddict-item]
+    )
+
+
+def check_json_v3(data: object) -> TypeGuard[AsTypeJSON_V3]:
+    """
+    A type guard for the Zarr V3 form of the Astype codec JSON
+    """
+    return (
+        _check_codecjson_v3(data)
+        and isinstance(data, Mapping)
+        and data["name"] == "astype"
+        and "configuration" in data
+        and "encode_dtype" in data["configuration"]
+        and "decode_dtype" in data["configuration"]
+        and check_dtype_spec_v3(data["configuration"]["decode_dtype"])
+        and check_dtype_spec_v3(data["configuration"]["encode_dtype"])
+    )
+
+
+class AsType(_NumcodecsArrayArrayCodec):
+    """
+    A wrapper around the numcodecs.Astype codec that provides Zarr V3 compatibility.
+
+    This class does not have a stable API.
+    """
+
+    codec_name = "numcodecs.astype"
+    _codec_id = "astype"
+    codec_config: AsTypeConfig_V2
+
+    @overload
+    def to_json(self, zarr_format: Literal[2]) -> AsTypeJSON_V2: ...
+    @overload
+    def to_json(self, zarr_format: Literal[3]) -> AsTypeJSON_V3: ...
+    def to_json(self, zarr_format: ZarrFormat) -> AsTypeJSON_V2 | AsTypeJSON_V3:
+        _warn_unstable_specification(self)
+        return super().to_json(zarr_format)  # type: ignore[return-value]
+
+    def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec:
+        dtype = parse_dtype(self.codec_config["encode_dtype"], zarr_format=3)
+        return replace(chunk_spec, dtype=dtype)
+
+    def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self:
+        dtype = array_spec.dtype.to_native_dtype()  # pragma: no cover
+        return type(self)(**{**self.codec_config, "decode_dtype": str(dtype)})  # pragma: no cover
+
+    @classmethod
+    def _from_json_v2(cls, data: CodecJSON_V2) -> Self:
+        return cls(**data)
+
+    @classmethod
+    def _from_json_v3(cls, data: CodecJSON_V3) -> Self:
+        if check_json_v3(data):
+            config = data["configuration"]
+            encode_dtype = parse_dtype(config["encode_dtype"], zarr_format=3).to_json(
+                zarr_format=2
+            )["name"]
+            decode_dtype = parse_dtype(config["decode_dtype"], zarr_format=3).to_json(
+                zarr_format=2
+            )["name"]
+
+            return cls(encode_dtype=encode_dtype, decode_dtype=decode_dtype)
+        raise TypeError(f"Invalid JSON: {data}")
+
+    @classmethod
+    def from_json(cls, data: CodecJSON) -> Self:
+        if _check_codecjson_v2(data):
+            return cls._from_json_v2(data)
+        return cls._from_json_v3(data)