[do not merge] proof of concept for unified v2 / v3 codecs

examples/image_codecs.py

@@ -0,0 +1,51 @@
+# /// script
+# requires-python = ">=3.11"
+# dependencies = [
+#   "zarr @ git+https://github.com/d-v-b/zarr-python.git@a2bc6555",
+#   "imagecodecs==2025.3.30",
+#   "pytest"
+# ]
+# ///
+
+#   "zarr @ git+https://github.com/zarr-developers/zarr-python.git@main",
+from typing import Literal
+
+import numcodecs
+import numpy as np
+import pytest
+from imagecodecs.numcodecs import Jpeg
+
+import zarr
+
+numcodecs.register_codec(Jpeg)
+jpg_codec = Jpeg()
+
+
+@pytest.mark.parametrize("zarr_format", [2, 3])
+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

@@ -82,7 +82,8 @@ test = [
     "pytest-xdist",
     "packaging",
     "tomlkit",
-    "uv"
+    "uv",
+    "imagecodecs"
 ]
 remote_tests = [
     'zarr[remote]',
@@ -383,6 +384,7 @@ module = [
     "tests.test_indexing",
     "tests.test_properties",
     "tests.test_sync",
+    "tests.test_v2",
     "tests.test_regression.scripts.*"
 ]
 ignore_errors = true

src/zarr/abc/codec.py

@@ -1,11 +1,21 @@
 from __future__ import annotations
 
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Generic, TypeVar
+from collections.abc import Mapping
+from typing import (
+    TYPE_CHECKING,
+    Generic,
+    Literal,
+    TypedDict,
+    TypeVar,
+    overload,
+)
+
+from typing_extensions import ReadOnly
 
 from zarr.abc.metadata import Metadata
 from zarr.core.buffer import Buffer, NDBuffer
-from zarr.core.common import ChunkCoords, concurrent_map
+from zarr.core.common import ChunkCoords, NamedConfig, ZarrFormat, concurrent_map
 from zarr.core.config import config
 
 if TYPE_CHECKING:
@@ -34,6 +44,21 @@
 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]):
+    id: ReadOnly[TName]
+
+
+CodecConfig_V3 = NamedConfig[str, Mapping[str, object]]
+
+CodecJSON_V3 = str | CodecConfig_V3
+
+# The widest type we will accept for a codec JSON
+# This covers v2 and v3
+CodecJSON = str | Mapping[str, object]
+
 
 class BaseCodec(Metadata, Generic[CodecInput, CodecOutput]):
     """Generic base class for codecs.
@@ -157,6 +182,34 @@ 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]: ...
+    @overload
+    def to_json(self, zarr_format: Literal[3]) -> NamedConfig[str, Mapping[str, object]]: ...
+
+    def to_json(
+        self, zarr_format: ZarrFormat
+    ) -> CodecJSON_V2[str] | NamedConfig[str, Mapping[str, object]]:
+        raise NotImplementedError
+
+    @classmethod
+    def _from_json_v2(cls, data: CodecJSON) -> Self:
+        raise NotImplementedError
+
+    @classmethod
+    def _from_json_v3(cls, data: CodecJSON) -> Self:
+        raise NotImplementedError
+
+    @classmethod
+    def from_json(cls, data: CodecJSON, zarr_format: ZarrFormat) -> Self:
+        if zarr_format == 2:
+            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
+
 
 class ArrayArrayCodec(BaseCodec[NDBuffer, NDBuffer]):
     """Base class for array-to-array codecs."""
@@ -447,3 +500,7 @@ async def wrap(chunk: CodecInput | None, chunk_spec: ArraySpec) -> CodecOutput |
         return await func(chunk, chunk_spec)
 
     return wrap
+
+
+# Raised when a codec JSON data is invalid
+class CodecValidationError(ValueError): ...

src/zarr/codecs/_v2.py

@@ -19,7 +19,7 @@
 
 
 @dataclass(frozen=True)
-class V2Codec(ArrayBytesCodec):
+class _V2Codec(ArrayBytesCodec):
     filters: tuple[numcodecs.abc.Codec, ...] | None
     compressor: numcodecs.abc.Codec | None