Seperate chunk key encoding 'name' (key) from

its implementation indicated by qualname.

Set default chunk key encodings implementations for `default` and `v2`
in global config.

Author: RFLeijenaar

src/zarr/core/chunk_key_encodings.py

@@ -108,5 +108,5 @@ def parse_chunk_key_encoding(data: ChunkKeyEncodingLike) -> ChunkKeyEncoding:
     return chunk_key_encoding
 
 
-register_chunk_key_encoding(DefaultChunkKeyEncoding, qualname="default")
-register_chunk_key_encoding(V2ChunkKeyEncoding, qualname="v2")
+register_chunk_key_encoding("default", DefaultChunkKeyEncoding)
+register_chunk_key_encoding("v2", V2ChunkKeyEncoding)

src/zarr/core/config.py

@@ -126,6 +126,10 @@ def enable_gpu(self) -> ConfigSet:
                 "vlen-utf8": "zarr.codecs.vlen_utf8.VLenUTF8Codec",
                 "vlen-bytes": "zarr.codecs.vlen_utf8.VLenBytesCodec",
             },
+            "chunk_key_encodings": {
+                "default": "zarr.core.chunk_key_encodings.DefaultChunkKeyEncoding",
+                "v2": "zarr.core.chunk_key_encodings.V2ChunkKeyEncoding",
+            },
             "buffer": "zarr.buffer.cpu.Buffer",
             "ndbuffer": "zarr.buffer.cpu.NDBuffer",
         }

src/zarr/registry.py

@@ -64,7 +64,8 @@ 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()
+# Now a dict[str, Registry[ChunkKeyEncoding]]
+__chunk_key_encoding_registries: dict[str, Registry[ChunkKeyEncoding]] = defaultdict(Registry)
 
 # CHANGE: Consider updating docstring
 """
@@ -105,12 +106,10 @@ 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")
-    )
+    for e in entry_points.select(group="zarr.chunk_key_encoding"):
+        __chunk_key_encoding_registries[e.name].lazy_load_list.append(e)
+    for e in entry_points.select(group="zarr", name="chunk_key_encoding"):
+        __chunk_key_encoding_registries[e.name].lazy_load_list.append(e)
 
     __pipeline_registry.lazy_load_list.extend(entry_points.select(group="zarr.codec_pipeline"))
     __pipeline_registry.lazy_load_list.extend(
@@ -127,7 +126,7 @@ def _collect_entrypoints() -> list[Registry[Any]]:
         __pipeline_registry,
         __buffer_registry,
         __ndbuffer_registry,
-        __chunk_key_encoding_registry,
+        *(__chunk_key_encoding_registries.values()),
     ]
 
 
@@ -158,8 +157,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 register_chunk_key_encoding(
+    key: str, cke_cls: type[ChunkKeyEncoding], qualname: str | None = None
+) -> None:
+    __chunk_key_encoding_registries[key].register(cke_cls, qualname)
 
 
 def get_codec_class(key: str, reload_config: bool = False) -> type[Codec]:
@@ -299,14 +300,36 @@ 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:
+def get_chunk_key_encoding_class(key: str, reload_config: bool = False) -> type[ChunkKeyEncoding]:
+    if reload_config:
+        _reload_config()
+
+    if key in __chunk_key_encoding_registries:
+        __chunk_key_encoding_registries[key].lazy_load()
+    else:
         raise KeyError(
-            f"Chunk key encoding '{key}' not found in registered chunk key encodings: {list(__chunk_key_encoding_registry)}."
+            f"Chunk key encoding '{key}' not found in registered chunk key encodings: {list(__chunk_key_encoding_registries)}."
         )
 
-    return __chunk_key_encoding_registry[key]
+    cke_classes = __chunk_key_encoding_registries[key]
+    if not cke_classes:
+        raise KeyError(key)
+
+    config_entry = config.get("chunk_key_encodings", {}).get(key)
+    if config_entry is None:
+        if len(cke_classes) == 1:
+            return next(iter(cke_classes.values()))
+        warnings.warn(
+            f"Chunk key encoding '{key}' not configured in config. Selecting any implementation.",
+            stacklevel=2,
+            category=ZarrUserWarning,
+        )
+        return list(cke_classes.values())[-1]
+    selected_encoding_cls = cke_classes[config_entry]
+
+    if selected_encoding_cls:
+        return selected_encoding_cls
+    raise KeyError(key)
 
 
 _collect_entrypoints()