Add v2 and v3 metadata support to codecs

.pre-commit-config.yaml

@@ -33,7 +33,7 @@ repos:
           - donfig
           - numcodecs[crc32c]
           - numpy==2.1 # until https://github.com/numpy/numpy/issues/28034 is resolved
-          - typing_extensions
+          - typing_extensions>=4.10.0
           - universal-pathlib
           - obstore>=0.5.1
           # Tests

changes/3332.feature.md

@@ -0,0 +1,5 @@
+Adds methods to the base `Codec` for creating Zarr V2 or Zarr V3 metadata for a given codec. This removes the distinction between Zarr V2 and Zarr V3 codecs, which will greatly simplify the conversion of Zarr V2 data to Zarr V3.
+
+This PR also changes how codecs defined in Numcodecs are handled: When a user attempts to create an array with a numcodecs codec (say, `numcodecs.Blosc`), Zarr will internally replace this codec with the `blosc` codec defined in Zarr. This ensures that Zarr controls the JSON metadata of the codecs we use, and that codecs are uniformly handled internally.
+
+Finally, this PR adds support for arbitrary codecs that implement the Numcodecs API. This codecs are wrapped in an adapter layer that imbues them with Zarr V3 codec properties. This adaptation is made on a best-effort basis, and in many cases may not be a substitute for writing a complete Zarr V3 codec.

examples/image_codec.py

@@ -0,0 +1,51 @@
+# /// script
+# requires-python = ">=3.11"
+# dependencies = [
+#   "zarr @ git+https://github.com/zarr-developers/zarr-python.git@main",
+#   "imagecodecs==2025.3.30",
+#   "pytest==8.4.2"
+# ]
+# ///
+
+from typing import Literal
+
+import numcodecs
+import numpy as np
+import pytest
+from imagecodecs.numcodecs import Jpeg
+
+import zarr
+from zarr.core.common import ZarrFormat
+
+numcodecs.register_codec(Jpeg)
+jpg_codec = Jpeg()
+
+
+@pytest.mark.parametrize("zarr_format", ZarrFormat)
+def test(zarr_format: Literal[2, 3]) -> None:
+    store = {}
+    if zarr_format == 2:
+        z_w = zarr.create_array(
+            store=store,
+            data=np.zeros((100, 100, 3), dtype=np.uint8),
+            compressors=jpg_codec,
+            zarr_format=zarr_format,
+        )
+    else:
+        z_w = zarr.create_array(
+            store=store,
+            data=np.zeros((100, 100, 3), dtype=np.uint8),
+            serializer=jpg_codec,
+            zarr_format=zarr_format,
+        )
+    z_w[:] = 2
+    z_r = zarr.open_array(store=store, zarr_format=zarr_format)
+    assert np.all(z_r[:] == 2)
+    if zarr_format == 2:
+        print(z_r.metadata.to_dict()["compressor"])
+    else:
+        print(z_r.metadata.to_dict()["codecs"])
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, f"-c {__file__}", "-s"])

pyproject.toml

@@ -35,7 +35,7 @@ dependencies = [
     'packaging>=22.0',
     'numpy>=1.26',
     'numcodecs[crc32c]>=0.14',
-    'typing_extensions>=4.9',
+    'typing_extensions>=4.10.0',
     'donfig>=0.8',
 ]
 
@@ -239,7 +239,7 @@ dependencies = [
     'fsspec==2023.10.0',
     's3fs==2023.10.0',
     'universal_pathlib==0.0.22',
-    'typing_extensions==4.9.*',
+    'typing_extensions>=4.10.0',
     'donfig==0.8.*',
     'obstore==0.5.*',
     # test deps

src/zarr/abc/codec.py

@@ -1,14 +1,25 @@
 from __future__ import annotations
 
 from abc import abstractmethod
-from collections.abc import Mapping
-from typing import TYPE_CHECKING, Generic, TypeGuard, TypeVar
-
-from typing_extensions import ReadOnly, TypedDict
+from typing import (
+    TYPE_CHECKING,
+    Generic,
+    Literal,
+    Self,
+    TypeVar,
+    overload,
+)
 
 from zarr.abc.metadata import Metadata
 from zarr.core.buffer import Buffer, NDBuffer
-from zarr.core.common import NamedConfig, concurrent_map
+from zarr.core.common import (
+    CodecJSON,
+    CodecJSON_V2,
+    CodecJSON_V3,
+    ZarrFormat,
+    _check_codecjson_v2,
+    concurrent_map,
+)
 from zarr.core.config import config
 
 if TYPE_CHECKING:
@@ -37,27 +48,6 @@
 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.
@@ -181,6 +171,30 @@ async def encode(
         """
         return await _batching_helper(self._encode_single, chunks_and_specs)
 
+    @overload
+    def to_json(self, zarr_format: Literal[2]) -> CodecJSON_V2: ...
+    @overload
+    def to_json(self, zarr_format: Literal[3]) -> CodecJSON_V3: ...
+
+    def to_json(self, zarr_format: ZarrFormat) -> CodecJSON_V2 | CodecJSON_V3:
+        raise NotImplementedError
+
+    @classmethod
+    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.get("configuration", {}))
+
+    @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)
+
 
 class ArrayArrayCodec(BaseCodec[NDBuffer, NDBuffer]):
     """Base class for array-to-array codecs."""

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

@@ -417,6 +417,7 @@ async def getsize(self, key: str) -> int:
         FileNotFoundError
             When the given key does not exist in the store.
         """
+
         # Note to implementers: this default implementation is very inefficient since
         # it requires reading the entire object. Many systems will have ways to get the
         # size of an object without reading it.

src/zarr/api/asynchronous.py

@@ -56,7 +56,6 @@
     from zarr.abc.numcodec import Numcodec
     from zarr.core.buffer import NDArrayLikeOrScalar
     from zarr.core.chunk_key_encodings import ChunkKeyEncoding
-    from zarr.core.metadata.v2 import CompressorLikev2
     from zarr.storage import StoreLike
 
     # TODO: this type could use some more thought
@@ -130,8 +129,8 @@ class _LikeArgs(TypedDict):
     chunks: NotRequired[tuple[int, ...]]
     dtype: NotRequired[np.dtype[np.generic]]
     order: NotRequired[Literal["C", "F"]]
-    filters: NotRequired[tuple[Numcodec, ...] | None]
-    compressor: NotRequired[CompressorLikev2]
+    filters: NotRequired[tuple[Codec, ...] | None]
+    compressor: NotRequired[Codec | None]
     codecs: NotRequired[tuple[Codec, ...]]
 
 

src/zarr/codecs/__init__.py

@@ -1,32 +1,28 @@
 from __future__ import annotations
 
-from zarr.codecs.blosc import BloscCname, BloscCodec, BloscShuffle
+from zarr.codecs.blosc import BloscCname, BloscCname_Lit, BloscCodec, BloscShuffle, BloscShuffle_Lit
 from zarr.codecs.bytes import BytesCodec, Endian
 from zarr.codecs.crc32c_ import Crc32cCodec
 from zarr.codecs.gzip import GzipCodec
 from zarr.codecs.numcodecs import (
     BZ2,
     CRC32,
-    CRC32C,
     LZ4,
     LZMA,
     ZFPY,
     Adler32,
     AsType,
     BitRound,
-    Blosc,
-    Delta,
     FixedScaleOffset,
     Fletcher32,
-    GZip,
     JenkinsLookup3,
     PackBits,
     PCodec,
     Quantize,
     Shuffle,
     Zlib,
-    Zstd,
 )
+from zarr.codecs.numcodecs.delta import Delta
 from zarr.codecs.sharding import ShardingCodec, ShardingCodecIndexLocation
 from zarr.codecs.transpose import TransposeCodec
 from zarr.codecs.vlen_utf8 import VLenBytesCodec, VLenUTF8Codec
@@ -35,8 +31,10 @@
 
 __all__ = [
     "BloscCname",
+    "BloscCname_Lit",
     "BloscCodec",
     "BloscShuffle",
+    "BloscShuffle_Lit",
     "BytesCodec",
     "Crc32cCodec",
     "Endian",
@@ -65,29 +63,62 @@
 # Register all the codecs formerly contained in numcodecs.zarr3
 
 register_codec("numcodecs.bz2", BZ2, qualname="zarr.codecs.numcodecs.BZ2")
+register_codec("bz2", BZ2, qualname="zarr.codecs.numcodecs.BZ2")
+
 register_codec("numcodecs.crc32", CRC32, qualname="zarr.codecs.numcodecs.CRC32")
-register_codec("numcodecs.crc32c", CRC32C, qualname="zarr.codecs.numcodecs.CRC32C")
+register_codec("crc32", CRC32, qualname="zarr.codecs.numcodecs.CRC32")
+
 register_codec("numcodecs.lz4", LZ4, qualname="zarr.codecs.numcodecs.LZ4")
+register_codec("lz4", LZ4, qualname="zarr.codecs.numcodecs.LZ4")
+
 register_codec("numcodecs.lzma", LZMA, qualname="zarr.codecs.numcodecs.LZMA")
+register_codec("lzma", LZMA, qualname="zarr.codecs.numcodecs.LZMA")
+
 register_codec("numcodecs.zfpy", ZFPY, qualname="zarr.codecs.numcodecs.ZFPY")
+register_codec("zfpy", ZFPY, qualname="zarr.codecs.numcodecs.ZFPY")
+
 register_codec("numcodecs.adler32", Adler32, qualname="zarr.codecs.numcodecs.Adler32")
+register_codec("adler32", Adler32, qualname="zarr.codecs.numcodecs.Adler32")
+
 register_codec("numcodecs.astype", AsType, qualname="zarr.codecs.numcodecs.AsType")
+register_codec("astype", AsType, qualname="zarr.codecs.numcodecs.AsType")
+
 register_codec("numcodecs.bitround", BitRound, qualname="zarr.codecs.numcodecs.BitRound")
-register_codec("numcodecs.blosc", Blosc, qualname="zarr.codecs.numcodecs.Blosc")
+register_codec("bitround", BitRound, qualname="zarr.codecs.numcodecs.BitRound")
+
 register_codec("numcodecs.delta", Delta, qualname="zarr.codecs.numcodecs.Delta")
+register_codec("delta", Delta, qualname="zarr.codecs.numcodecs.Delta")
+
 register_codec(
     "numcodecs.fixedscaleoffset",
     FixedScaleOffset,
     qualname="zarr.codecs.numcodecs.FixedScaleOffset",
 )
+register_codec(
+    "fixedscaleoffset",
+    FixedScaleOffset,
+    qualname="zarr.codecs.numcodecs.FixedScaleOffset",
+)
+
 register_codec("numcodecs.fletcher32", Fletcher32, qualname="zarr.codecs.numcodecs.Fletcher32")
-register_codec("numcodecs.gzip", GZip, qualname="zarr.codecs.numcodecs.GZip")
+register_codec("fletcher32", Fletcher32, qualname="zarr.codecs.numcodecs.Fletcher32")
+
 register_codec(
     "numcodecs.jenkins_lookup3", JenkinsLookup3, qualname="zarr.codecs.numcodecs.JenkinsLookup3"
 )
+register_codec("jenkins_lookup3", JenkinsLookup3, qualname="zarr.codecs.numcodecs.JenkinsLookup3")
+
+register_codec("pcodec", PCodec, qualname="zarr.codecs.numcodecs.PCodec")
 register_codec("numcodecs.pcodec", PCodec, qualname="zarr.codecs.numcodecs.PCodec")
+
 register_codec("numcodecs.packbits", PackBits, qualname="zarr.codecs.numcodecs.PackBits")
+register_codec("packbits", PackBits, qualname="zarr.codecs.numcodecs.PackBits")
+
 register_codec("numcodecs.quantize", Quantize, qualname="zarr.codecs.numcodecs.Quantize")
+register_codec("quantize", Quantize, qualname="zarr.codecs.numcodecs.Quantize")
+
 register_codec("numcodecs.shuffle", Shuffle, qualname="zarr.codecs.numcodecs.Shuffle")
+register_codec("shuffle", Shuffle, qualname="zarr.codecs.numcodecs.Shuffle")
+
 register_codec("numcodecs.zlib", Zlib, qualname="zarr.codecs.numcodecs.Zlib")
-register_codec("numcodecs.zstd", Zstd, qualname="zarr.codecs.numcodecs.Zstd")
+register_codec("zlib", Zlib, qualname="zarr.codecs.numcodecs.Zlib")