Use ChunkKeyEncodingLike type alias (zarr-developers#2763)

* make proper ChunkKeyEncodingLike type alias, and use it

* model the not-required-ness of the separator

* add cast

* changelog

Author: d-v-b

changes/2763.chore.rst

@@ -0,0 +1,3 @@
+Created a type alias ``ChunkKeyEncodingLike`` to model the union of ``ChunkKeyEncoding`` instances and the dict form of the
+parameters of those instances. ``ChunkKeyEncodingLike`` should be used by high-level functions to provide a convenient
+way for creating ``ChunkKeyEncoding`` objects.

src/zarr/core/array.py

@@ -412,7 +412,7 @@ async def create(
         # v3 only
         chunk_shape: ShapeLike | None = None,
         chunk_key_encoding: (
-            ChunkKeyEncoding
+            ChunkKeyEncodingLike
             | tuple[Literal["default"], Literal[".", "/"]]
             | tuple[Literal["v2"], Literal[".", "/"]]
             | None
@@ -453,7 +453,7 @@ async def create(
             The shape of the array's chunks
             Zarr format 3 only. Zarr format 2 arrays should use `chunks` instead.
             If not specified, default are guessed based on the shape and dtype.
-        chunk_key_encoding : ChunkKeyEncoding, optional
+        chunk_key_encoding : ChunkKeyEncodingLike, optional
             A specification of how the chunk keys are represented in storage.
             Zarr format 3 only. Zarr format 2 arrays should use `dimension_separator` instead.
             Default is ``("default", "/")``.
@@ -553,7 +553,7 @@ async def _create(
         # v3 only
         chunk_shape: ShapeLike | None = None,
         chunk_key_encoding: (
-            ChunkKeyEncoding
+            ChunkKeyEncodingLike
             | tuple[Literal["default"], Literal[".", "/"]]
             | tuple[Literal["v2"], Literal[".", "/"]]
             | None
@@ -671,7 +671,7 @@ async def _create_v3(
         config: ArrayConfig,
         fill_value: Any | None = None,
         chunk_key_encoding: (
-            ChunkKeyEncoding
+            ChunkKeyEncodingLike
             | tuple[Literal["default"], Literal[".", "/"]]
             | tuple[Literal["v2"], Literal[".", "/"]]
             | None
@@ -1708,7 +1708,7 @@ def create(
             The shape of the Array's chunks.
             Zarr format 3 only. Zarr format 2 arrays should use `chunks` instead.
             If not specified, default are guessed based on the shape and dtype.
-        chunk_key_encoding : ChunkKeyEncoding, optional
+        chunk_key_encoding : ChunkKeyEncodingLike, optional
             A specification of how the chunk keys are represented in storage.
             Zarr format 3 only. Zarr format 2 arrays should use `dimension_separator` instead.
             Default is ``("default", "/")``.
@@ -3756,7 +3756,7 @@ async def create_array(
     order: MemoryOrder | None = None,
     zarr_format: ZarrFormat | None = 3,
     attributes: dict[str, JSON] | None = None,
-    chunk_key_encoding: ChunkKeyEncoding | ChunkKeyEncodingLike | None = None,
+    chunk_key_encoding: ChunkKeyEncodingLike | None = None,
     dimension_names: Iterable[str] | None = None,
     storage_options: dict[str, Any] | None = None,
     overwrite: bool = False,
@@ -3834,7 +3834,7 @@ async def create_array(
         The zarr format to use when saving.
     attributes : dict, optional
         Attributes for the array.
-    chunk_key_encoding : ChunkKeyEncoding, optional
+    chunk_key_encoding : ChunkKeyEncodingLike, optional
         A specification of how the chunk keys are represented in storage.
         For Zarr format 3, the default is ``{"name": "default", "separator": "/"}}``.
         For Zarr format 2, the default is ``{"name": "v2", "separator": "."}}``.

src/zarr/core/chunk_key_encodings.py

@@ -2,7 +2,10 @@
 
 from abc import abstractmethod
 from dataclasses import dataclass
-from typing import Literal, TypedDict, cast
+from typing import TYPE_CHECKING, Literal, TypeAlias, TypedDict, cast
+
+if TYPE_CHECKING:
+    from typing import NotRequired
 
 from zarr.abc.metadata import Metadata
 from zarr.core.common import (
@@ -20,9 +23,9 @@ def parse_separator(data: JSON) -> SeparatorLiteral:
     return cast(SeparatorLiteral, data)
 
 
-class ChunkKeyEncodingLike(TypedDict):
+class ChunkKeyEncodingParams(TypedDict):
     name: Literal["v2", "default"]
-    separator: SeparatorLiteral
+    separator: NotRequired[SeparatorLiteral]
 
 
 @dataclass(frozen=True)
@@ -36,16 +39,17 @@ def __init__(self, *, separator: SeparatorLiteral) -> None:
         object.__setattr__(self, "separator", separator_parsed)
 
     @classmethod
-    def from_dict(
-        cls, data: dict[str, JSON] | ChunkKeyEncoding | ChunkKeyEncodingLike
-    ) -> ChunkKeyEncoding:
+    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
         name_parsed, config_parsed = parse_named_configuration(data, require_configuration=False)
         if name_parsed == "default":
@@ -73,6 +77,9 @@ def encode_chunk_key(self, chunk_coords: ChunkCoords) -> str:
         pass
 
 
+ChunkKeyEncodingLike: TypeAlias = ChunkKeyEncodingParams | ChunkKeyEncoding
+
+
 @dataclass(frozen=True)
 class DefaultChunkKeyEncoding(ChunkKeyEncoding):
     name: Literal["default"] = "default"

src/zarr/core/metadata/v3.py

@@ -27,7 +27,7 @@
 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
+from zarr.core.chunk_key_encodings import ChunkKeyEncoding, ChunkKeyEncodingLike
 from zarr.core.common import (
     JSON,
     ZARR_JSON,
@@ -253,7 +253,7 @@ def __init__(
         shape: Iterable[int],
         data_type: npt.DTypeLike | DataType,
         chunk_grid: dict[str, JSON] | ChunkGrid,
-        chunk_key_encoding: dict[str, JSON] | ChunkKeyEncoding,
+        chunk_key_encoding: ChunkKeyEncodingLike,
         fill_value: Any,
         codecs: Iterable[Codec | dict[str, JSON]],
         attributes: dict[str, JSON] | None,