refactor codec parsing and add tests

Author: d-v-b

src/zarr/codecs/_v2.py

@@ -15,7 +15,6 @@
 from zarr.core.common import (
     CodecJSON,
     CodecJSON_V2,
-    CodecJSON_V3,
     _check_codecjson_v2,
     _check_codecjson_v3,
 )
@@ -32,26 +31,6 @@
     from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar, ZDType
 
 
-def codec_json_v2_to_v3(data: CodecJSON_V2) -> CodecJSON_V3:
-    """
-    Convert V2 codec JSON to V3 codec JSON
-    """
-    name = data["id"]
-    config = {k: v for k, v in data.items() if k != "id"}
-    return {"name": name, "configuration": config}
-
-
-def codec_json_v3_to_v2(data: CodecJSON_V3) -> CodecJSON_V2:
-    """
-    Convert V3 codec JSON to V2 codec JSON
-    """
-    if isinstance(data, str):
-        return {"id": data}
-    name = data["name"]
-    config = dict(data.get("configuration", {}))
-    return {"id": name, **config}  # type: ignore[typeddict-item]
-
-
 @dataclass(frozen=True)
 class V2Codec(ArrayBytesCodec):
     filters: tuple[Numcodec, ...] | None
@@ -168,8 +147,12 @@ def _from_json_v2(cls, data: CodecJSON) -> Self:
     @classmethod
     def _from_json_v3(cls, data: CodecJSON) -> Self:
         if _check_codecjson_v3(data):
-            # convert to a v2 codec JSON
-            codec = get_numcodec(codec_json_v3_to_v2(data))
+            request: CodecJSON_V2
+            if isinstance(data, str):
+                request = {"id": data}
+            else:
+                request = {"id": data["name"], **data["configuration"]}  # type: ignore[typeddict-item]
+            codec = get_numcodec(request)
             return cls(codec=codec)
         msg = (
             "Invalid Zarr V3 JSON representation of a codec. "

src/zarr/codecs/numcodecs/_codecs.py

@@ -75,6 +75,7 @@
     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
@@ -154,21 +155,32 @@ class ShuffleJSON_V3(NamedRequiredConfig[Literal["shuffle"], ShuffleConfig]):
     """JSON representation of Shuffle codec for Zarr V3."""
 
 
-# Array-to-array codec configuration classes
-class DeltaConfig(TypedDict):
-    dtype: str
-    astype: str
+class DeltaConfig_V2(TypedDict):
+    dtype: DTypeSpec_V2
+    astype: DTypeSpec_V2
+
+
+class DeltaConfig_V3(TypedDict):
+    dtype: DTypeSpec_V3
+    astype: DTypeSpec_V3
 
 
 class BitRoundConfig(TypedDict):
     keepbits: int
 
 
-class FixedScaleOffsetConfig(TypedDict):
-    dtype: NotRequired[str]
+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]
-    astype: NotRequired[str]
 
 
 class QuantizeConfig(TypedDict):
@@ -186,13 +198,13 @@ class AsTypeConfig(TypedDict):
 
 
 # Array-to-array codec JSON representations
-class DeltaJSON_V2(DeltaConfig):
+class DeltaJSON_V2(DeltaConfig_V2):
     """JSON representation of Delta codec for Zarr V2."""
 
     id: ReadOnly[Literal["delta"]]
 
 
-class DeltaJSON_V3(NamedRequiredConfig[Literal["delta"], DeltaConfig]):
+class DeltaJSON_V3(NamedRequiredConfig[Literal["delta"], DeltaConfig_V3]):
     """JSON representation of Delta codec for Zarr V3."""
 
 
@@ -206,14 +218,14 @@ class BitRoundJSON_V3(NamedRequiredConfig[Literal["bitround"], BitRoundConfig]):
     """JSON representation of BitRound codec for Zarr V3."""
 
 
-class FixedScaleOffsetJSON_V2(FixedScaleOffsetConfig):
+class FixedScaleOffsetJSON_V2(FixedScaleOffsetConfig_V2):
     """JSON representation of FixedScaleOffset codec for Zarr V2."""
 
     id: ReadOnly[Literal["fixedscaleoffset"]]
 
 
 class FixedScaleOffsetJSON_V3(
-    NamedRequiredConfig[Literal["fixedscaleoffset"], FixedScaleOffsetConfig]
+    NamedRequiredConfig[Literal["fixedscaleoffset"], FixedScaleOffsetConfig_V3]
 ):
     """JSON representation of FixedScaleOffset codec for Zarr V3."""
 
@@ -507,7 +519,11 @@ def to_json(self, zarr_format: Literal[2]) -> ZstdJSON_V2: ...
     @overload
     def to_json(self, zarr_format: Literal[3]) -> ZstdJSON_V3: ...
     def to_json(self, zarr_format: ZarrFormat) -> ZstdJSON_V2 | ZstdJSON_V3:
-        return super().to_json(zarr_format)  # type: ignore[return-value]
+        res = super().to_json(zarr_format)
+        if zarr_format == 2 and not res.get("checksum", False):  # type: ignore[union-attr]
+            # https://github.com/zarr-developers/zarr-python/pull/2655
+            res.pop("checksum")  # type: ignore[union-attr, typeddict-item]
+        return res  # type: ignore[return-value]
 
 
 class Zlib(_NumcodecsBytesBytesCodec):
@@ -595,7 +611,7 @@ def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self:
 class Delta(_NumcodecsArrayArrayCodec):
     codec_name = "numcodecs.delta"
     _codec_id = "delta"
-    codec_config: DeltaConfig
+    codec_config: DeltaConfig_V2 | DeltaConfig_V3
 
     def __init__(self, **codec_config: Any) -> None:
         if "codec_config" in codec_config:
@@ -612,7 +628,7 @@ def to_json(self, zarr_format: ZarrFormat) -> DeltaJSON_V2 | DeltaJSON_V3:
 
     def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec:
         if astype := self.codec_config.get("astype"):
-            dtype = parse_dtype(np.dtype(astype), zarr_format=3)
+            dtype = parse_dtype(np.dtype(astype), zarr_format=3)  # type: ignore[arg-type]
             return replace(chunk_spec, dtype=dtype)
         return chunk_spec
 
@@ -634,7 +650,7 @@ def to_json(self, zarr_format: ZarrFormat) -> BitRoundJSON_V2 | BitRoundJSON_V3:
 class FixedScaleOffset(_NumcodecsArrayArrayCodec):
     codec_name = "numcodecs.fixedscaleoffset"
     _codec_id = "fixedscaleoffset"
-    codec_config: FixedScaleOffsetConfig
+    codec_config: FixedScaleOffsetConfig_V2
 
     @overload
     def to_json(self, zarr_format: Literal[2]) -> FixedScaleOffsetJSON_V2: ...
@@ -646,7 +662,7 @@ def to_json(self, zarr_format: ZarrFormat) -> FixedScaleOffsetJSON_V2 | FixedSca
 
     def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec:
         if astype := self.codec_config.get("astype"):
-            dtype = parse_dtype(np.dtype(astype), zarr_format=3)
+            dtype = parse_dtype(np.dtype(astype), zarr_format=3)  # type: ignore[arg-type]
             return replace(chunk_spec, dtype=dtype)
         return chunk_spec
 

src/zarr/codecs/zstd.py

@@ -4,7 +4,7 @@
 from collections.abc import Mapping
 from dataclasses import dataclass
 from functools import cached_property
-from typing import TYPE_CHECKING, Literal, Self, TypedDict, TypeGuard, cast, overload
+from typing import TYPE_CHECKING, Literal, NotRequired, Self, TypedDict, TypeGuard, cast, overload
 
 import numcodecs
 from numcodecs.zstd import Zstd
@@ -25,6 +25,7 @@
 
 class ZstdConfig_V2(TypedDict):
     level: int
+    checksum: NotRequired[Literal[True]]
 
 
 class ZstdConfig_V3(TypedDict):
@@ -104,7 +105,7 @@ def from_dict(cls, data: dict[str, JSON]) -> Self:
     def _from_json_v2(cls, data: CodecJSON) -> Self:
         if check_json_v2(data):
             if "checksum" in data:
-                return cls(level=data["level"], checksum=data["checksum"])  # type: ignore[typeddict-item]
+                return cls(level=data["level"], checksum=data["checksum"])
             else:
                 return cls(level=data["level"])
 
@@ -138,7 +139,10 @@ def to_json(self, zarr_format: Literal[3]) -> ZstdJSON_V3: ...
 
     def to_json(self, zarr_format: ZarrFormat) -> ZstdJSON_V2 | ZstdJSON_V3:
         if zarr_format == 2:
-            return {"id": "zstd", "level": self.level}
+            if self.checksum is True:
+                return {"id": "zstd", "level": self.level, "checksum": self.checksum}
+            else:
+                return {"id": "zstd", "level": self.level}
         else:
             return {
                 "name": "zstd",

src/zarr/core/array.py

@@ -25,6 +25,7 @@
 import zarr
 from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BytesBytesCodec, Codec
 from zarr.abc.numcodec import Numcodec
+from zarr.codecs._v2 import NumcodecWrapper
 from zarr.codecs.bytes import BytesCodec
 from zarr.codecs.transpose import TransposeCodec
 from zarr.codecs.vlen_utf8 import VLenBytesCodec, VLenUTF8Codec
@@ -109,6 +110,7 @@
     ArrayV3MetadataDict,
     T_ArrayMetadata,
 )
+from zarr.core.metadata.common import _parse_codec
 from zarr.core.metadata.io import save_metadata
 from zarr.core.metadata.v2 import (
     CompressorLike_V2,
@@ -125,9 +127,6 @@
     ZarrUserWarning,
 )
 from zarr.registry import (
-    _parse_array_array_codec,
-    _parse_array_bytes_codec,
-    _parse_bytes_bytes_codec,
     get_pipeline_class,
 )
 from zarr.storage._common import StorePath, ensure_no_existing_node, make_store_path
@@ -141,7 +140,6 @@
 
     from zarr.abc.codec import CodecPipeline
     from zarr.abc.store import Store
-    from zarr.codecs._v2 import NumcodecWrapper
     from zarr.codecs.sharding import ShardingCodecIndexLocation
     from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar
     from zarr.storage import StoreLike
@@ -252,6 +250,60 @@ def create_codec_pipeline(metadata: ArrayMetadata, *, store: Store | None = None
     raise TypeError  # pragma: no cover
 
 
+def _parse_bytes_bytes_codec(
+    data: Mapping[str, JSON] | Codec | Numcodec, *, dtype: ZDType[Any, Any]
+) -> BytesBytesCodec:
+    """
+    Normalize the input to a ``BytesBytesCodec`` instance.
+    If the input is already a ``BytesBytesCodec``, it is returned as is. If the input is a dict, it
+    is converted to a ``BytesBytesCodec`` instance via the ``_resolve_codec`` function.
+    """
+
+    _codec_or_numcodec_wrapper = _parse_codec(data, dtype=dtype)
+    if isinstance(_codec_or_numcodec_wrapper, NumcodecWrapper):
+        return _codec_or_numcodec_wrapper.to_bytes_bytes()
+    elif isinstance(_codec_or_numcodec_wrapper, BytesBytesCodec):
+        return _codec_or_numcodec_wrapper
+    msg = f"Expected a NumcodecWrapper or ArrayBytesCodec or a dict representation thereof; got {data} instead."
+    raise ValueError(msg)
+
+
+def _parse_array_bytes_codec(
+    data: Mapping[str, JSON] | Codec | Numcodec, *, dtype: ZDType[Any, Any]
+) -> ArrayBytesCodec:
+    """
+    Normalize the input to a ``ArrayBytesCodec`` instance.
+    If the input is already a ``ArrayBytesCodec``, it is returned as is. If the input is a dict, it
+    is converted to a ``ArrayBytesCodec`` instance via the ``_resolve_codec`` function.
+    """
+
+    _codec_or_numcodec_wrapper = _parse_codec(data, dtype=dtype)
+    if isinstance(_codec_or_numcodec_wrapper, NumcodecWrapper):
+        return _codec_or_numcodec_wrapper.to_array_bytes()
+    elif isinstance(_codec_or_numcodec_wrapper, ArrayBytesCodec):
+        return _codec_or_numcodec_wrapper
+    msg = f"Expected a NumcodecWrapper or ArrayBytesCodec or a dict representation thereof; got {data} instead."
+    raise ValueError(msg)
+
+
+def _parse_array_array_codec(
+    data: Mapping[str, JSON] | Codec | Numcodec, *, dtype: ZDType[Any, Any]
+) -> ArrayArrayCodec:
+    """
+    Normalize the input to a ``ArrayArrayCodec`` instance.
+    If the input is already a ``ArrayArrayCodec``, it is returned as is. If the input is a dict, it
+    is converted to a ``ArrayArrayCodec`` instance via the ``_resolve_codec`` function.
+    """
+
+    _codec_or_numcodec_wrapper = _parse_codec(data, dtype=dtype)
+    if isinstance(_codec_or_numcodec_wrapper, NumcodecWrapper):
+        return _codec_or_numcodec_wrapper.to_array_array()
+    elif isinstance(_codec_or_numcodec_wrapper, ArrayArrayCodec):
+        return _codec_or_numcodec_wrapper
+    msg = f"Expected a NumcodecWrapper or ArrayArrayCodec or a dict representation thereof; got {data} instead."
+    raise ValueError(msg)
+
+
 async def get_array_metadata(
     store_path: StorePath, zarr_format: ZarrFormat | None = 3
 ) -> dict[str, JSON]:
@@ -5177,17 +5229,15 @@ def _parse_chunk_encoding_v3(
             maybe_array_array = (filters,)
         else:
             maybe_array_array = cast("Iterable[Codec | dict[str, JSON]]", filters)
-        out_array_array = tuple(
-            _parse_array_array_codec(c, zarr_format=3) for c in maybe_array_array
-        )
+        out_array_array = tuple(_parse_array_array_codec(c, dtype=dtype) for c in maybe_array_array)
 
     if serializer == "auto":
         out_array_bytes = default_serializer_v3(dtype)
     else:
         # TODO: ensure that the serializer is compatible with the ndarray produced by the
         # array-array codecs. For example, if a sequence of array-array codecs produces an
         # array with a single-byte data type, then the serializer should not specify endiannesss.
-        out_array_bytes = _parse_array_bytes_codec(serializer, zarr_format=3)
+        out_array_bytes = _parse_array_bytes_codec(serializer, dtype=dtype)
 
     if compressors is None:
         out_bytes_bytes: tuple[BytesBytesCodec, ...] = ()
@@ -5200,12 +5250,7 @@ def _parse_chunk_encoding_v3(
         else:
             maybe_bytes_bytes = compressors  # type: ignore[assignment]
 
-        out_bytes_bytes = tuple(_parse_bytes_bytes_codec(c) for c in maybe_bytes_bytes)
-
-    # specialize codecs as needed given the dtype
-
-    # TODO: refactor so that the config only contains the name of the codec, and we use the dtype
-    # to create the codec instance, instead of storing a dict representation of a full codec.
+        out_bytes_bytes = tuple(_parse_bytes_bytes_codec(c, dtype=dtype) for c in maybe_bytes_bytes)
 
     # TODO: ensure that the serializer is compatible with the ndarray produced by the
     # array-array codecs. For example, if a sequence of array-array codecs produces an