Add registry for chunk key encodings.

Author: RFLeijenaar

src/zarr/core/array.py

@@ -47,6 +47,7 @@
     ChunkKeyEncodingLike,
     DefaultChunkKeyEncoding,
     V2ChunkKeyEncoding,
+    parse_chunk_key_encoding,
 )
 from zarr.core.common import (
     JSON,
@@ -4934,13 +4935,11 @@ def _parse_chunk_key_encoding(
     """
     if data is None:
         if zarr_format == 2:
-            result = ChunkKeyEncoding.from_dict({"name": "v2", "separator": "."})
+            data = {"name": "v2", "configuration": {"separator": "."}}
         else:
-            result = ChunkKeyEncoding.from_dict({"name": "default", "separator": "/"})
-    elif isinstance(data, ChunkKeyEncoding):
-        result = data
-    else:
-        result = ChunkKeyEncoding.from_dict(data)
+            data = {"name": "default", "configuration": {"separator": "/"}}
+    result = parse_chunk_key_encoding(data)
+
     if zarr_format == 2 and result.name != "v2":
         msg = (
             "Invalid chunk key encoding. For Zarr format 2 arrays, the `name` field of the "

src/zarr/core/chunk_key_encodings.py

@@ -5,13 +5,14 @@
 from typing import TYPE_CHECKING, Literal, TypeAlias, TypedDict, cast
 
 if TYPE_CHECKING:
-    from typing import NotRequired
+    from typing import NotRequired, Self
 
 from zarr.abc.metadata import Metadata
 from zarr.core.common import (
     JSON,
     parse_named_configuration,
 )
+from zarr.registry import get_chunk_key_encoding_class, register_chunk_key_encoding
 
 SeparatorLiteral = Literal[".", "/"]
 
@@ -38,31 +39,9 @@ def __init__(self, *, separator: SeparatorLiteral) -> None:
         object.__setattr__(self, "separator", separator_parsed)
 
     @classmethod
-    def from_dict(cls, data: dict[str, JSON] | ChunkKeyEncodingLike) -> ChunkKeyEncoding:
-        if isinstance(data, ChunkKeyEncoding):
-            return data
-
-        # handle ChunkKeyEncodingParams
-        if "name" in data and "separator" in data:
-            data = {"name": data["name"], "configuration": {"separator": data["separator"]}}
-
-        # TODO: remove this cast when we are statically typing the JSON metadata completely.
-        data = cast("dict[str, JSON]", data)
-
-        # configuration is optional for chunk key encodings
+    def from_dict(cls, data: dict[str, JSON]) -> Self:
         name_parsed, config_parsed = parse_named_configuration(data, require_configuration=False)
-        if name_parsed == "default":
-            if config_parsed is None:
-                # for default, normalize missing configuration to use the "/" separator.
-                config_parsed = {"separator": "/"}
-            return DefaultChunkKeyEncoding(**config_parsed)  # type: ignore[arg-type]
-        if name_parsed == "v2":
-            if config_parsed is None:
-                # for v2, normalize missing configuration to use the "." separator.
-                config_parsed = {"separator": "."}
-            return V2ChunkKeyEncoding(**config_parsed)  # type: ignore[arg-type]
-        msg = f"Unknown chunk key encoding. Got {name_parsed}, expected one of ('v2', 'default')."
-        raise ValueError(msg)
+        return cls(**config_parsed if config_parsed else {})  # type: ignore[arg-type]
 
     def to_dict(self) -> dict[str, JSON]:
         return {"name": self.name, "configuration": {"separator": self.separator}}
@@ -76,12 +55,13 @@ def encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str:
         pass
 
 
-ChunkKeyEncodingLike: TypeAlias = ChunkKeyEncodingParams | ChunkKeyEncoding
+ChunkKeyEncodingLike: TypeAlias = dict[str, JSON] | ChunkKeyEncodingParams | ChunkKeyEncoding
 
 
 @dataclass(frozen=True)
 class DefaultChunkKeyEncoding(ChunkKeyEncoding):
     name: Literal["default"] = "default"
+    separator: SeparatorLiteral = "/"  # default
 
     def decode_chunk_key(self, chunk_key: str) -> tuple[int, ...]:
         if chunk_key == "c":
@@ -95,10 +75,38 @@ def encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str:
 @dataclass(frozen=True)
 class V2ChunkKeyEncoding(ChunkKeyEncoding):
     name: Literal["v2"] = "v2"
+    separator: SeparatorLiteral = "."  # default
 
     def decode_chunk_key(self, chunk_key: str) -> tuple[int, ...]:
         return tuple(map(int, chunk_key.split(self.separator)))
 
     def encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str:
         chunk_identifier = self.separator.join(map(str, chunk_coords))
         return "0" if chunk_identifier == "" else chunk_identifier
+
+
+def parse_chunk_key_encoding(data: ChunkKeyEncodingLike) -> ChunkKeyEncoding:
+    """
+    Take an implicit specification of a chunk key encoding and parse it into a ChunkKeyEncoding object.
+    """
+    if isinstance(data, ChunkKeyEncoding):
+        return data
+
+    # handle ChunkKeyEncodingParams
+    if "name" in data and "separator" in data:
+        data = {"name": data["name"], "configuration": {"separator": data["separator"]}}
+
+    # Now must be a named config
+    data = cast("dict[str, JSON]", data)
+
+    name_parsed, _ = parse_named_configuration(data, require_configuration=False)
+    try:
+        chunk_key_encoding = get_chunk_key_encoding_class(name_parsed).from_dict(data)
+    except KeyError as e:
+        raise ValueError(f"Unknown chunk key encoding: {e.args[0]!r}") from e
+
+    return chunk_key_encoding
+
+
+register_chunk_key_encoding(DefaultChunkKeyEncoding, qualname="default")
+register_chunk_key_encoding(V2ChunkKeyEncoding, qualname="v2")

src/zarr/core/metadata/v3.py

@@ -24,7 +24,11 @@
 from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BytesBytesCodec, Codec
 from zarr.core.array_spec import ArrayConfig, ArraySpec
 from zarr.core.chunk_grids import ChunkGrid, RegularChunkGrid
-from zarr.core.chunk_key_encodings import ChunkKeyEncoding, ChunkKeyEncodingLike
+from zarr.core.chunk_key_encodings import (
+    ChunkKeyEncoding,
+    ChunkKeyEncodingLike,
+    parse_chunk_key_encoding,
+)
 from zarr.core.common import (
     JSON,
     ZARR_JSON,
@@ -174,7 +178,7 @@ def __init__(
 
         shape_parsed = parse_shapelike(shape)
         chunk_grid_parsed = ChunkGrid.from_dict(chunk_grid)
-        chunk_key_encoding_parsed = ChunkKeyEncoding.from_dict(chunk_key_encoding)
+        chunk_key_encoding_parsed = parse_chunk_key_encoding(chunk_key_encoding)
         dimension_names_parsed = parse_dimension_names(dimension_names)
         # Note: relying on a type method is numpy-specific
         fill_value_parsed = data_type.cast_scalar(fill_value)

src/zarr/registry.py

@@ -22,15 +22,19 @@
     )
     from zarr.abc.numcodec import Numcodec
     from zarr.core.buffer import Buffer, NDBuffer
+    from zarr.core.chunk_key_encodings import ChunkKeyEncoding
     from zarr.core.common import JSON
 
+# CHANGE: Consider adding here
 __all__ = [
     "Registry",
     "get_buffer_class",
+    "get_chunk_key_encoding_class",
     "get_codec_class",
     "get_ndbuffer_class",
     "get_pipeline_class",
     "register_buffer",
+    "register_chunk_key_encoding",
     "register_codec",
     "register_ndbuffer",
     "register_pipeline",
@@ -60,10 +64,12 @@ def register(self, cls: type[T], qualname: str | None = None) -> None:
 __pipeline_registry: Registry[CodecPipeline] = Registry()
 __buffer_registry: Registry[Buffer] = Registry()
 __ndbuffer_registry: Registry[NDBuffer] = Registry()
+__chunk_key_encoding_registry: Registry[ChunkKeyEncoding] = Registry()
 
+# CHANGE: Consider updating docstring
 """
 The registry module is responsible for managing implementations of codecs,
-pipelines, buffers and ndbuffers and collecting them from entrypoints.
+pipelines, buffers, ndbuffers, and chunk key encodings and collecting them from entrypoints.
 The implementation used is determined by the config.
 
 The registry module is also responsible for managing dtypes.
@@ -99,6 +105,13 @@ def _collect_entrypoints() -> list[Registry[Any]]:
     data_type_registry._lazy_load_list.extend(entry_points.select(group="zarr.data_type"))
     data_type_registry._lazy_load_list.extend(entry_points.select(group="zarr", name="data_type"))
 
+    __chunk_key_encoding_registry.lazy_load_list.extend(
+        entry_points.select(group="zarr.chunk_key_encoding")
+    )
+    __chunk_key_encoding_registry.lazy_load_list.extend(
+        entry_points.select(group="zarr", name="chunk_key_encoding")
+    )
+
     __pipeline_registry.lazy_load_list.extend(entry_points.select(group="zarr.codec_pipeline"))
     __pipeline_registry.lazy_load_list.extend(
         entry_points.select(group="zarr", name="codec_pipeline")
@@ -114,6 +127,7 @@ def _collect_entrypoints() -> list[Registry[Any]]:
         __pipeline_registry,
         __buffer_registry,
         __ndbuffer_registry,
+        __chunk_key_encoding_registry,
     ]
 
 
@@ -144,6 +158,10 @@ def register_buffer(cls: type[Buffer], qualname: str | None = None) -> None:
     __buffer_registry.register(cls, qualname)
 
 
+def register_chunk_key_encoding(cls: type, qualname: str | None = None) -> None:
+    __chunk_key_encoding_registry.register(cls, qualname)
+
+
 def get_codec_class(key: str, reload_config: bool = False) -> type[Codec]:
     if reload_config:
         _reload_config()
@@ -281,6 +299,16 @@ def get_ndbuffer_class(reload_config: bool = False) -> type[NDBuffer]:
     )
 
 
+def get_chunk_key_encoding_class(key: str) -> type[ChunkKeyEncoding]:
+    __chunk_key_encoding_registry.lazy_load()
+    if key not in __chunk_key_encoding_registry:
+        raise KeyError(
+            f"Chunk key encoding '{key}' not found in registered chunk key encodings: {list(__chunk_key_encoding_registry)}."
+        )
+
+    return __chunk_key_encoding_registry[key]
+
+
 _collect_entrypoints()