Return TypedDict instances from Metadata.to_dict()

Define TypedDict classes for metadata models that have a well-typed dict
representation and set the return value of the the model's to_dict() method
to the TypedDict

Author: ahmadjiha

src/zarr/abc/codec.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Any, Generic, TypeVar
+from typing import TYPE_CHECKING, Any, Generic, NotRequired, TypedDict, TypeVar
 
 from zarr.abc.metadata import Metadata
 from zarr.core.buffer import Buffer, NDBuffer
@@ -174,6 +174,22 @@ class BytesBytesCodec(BaseCodec[Buffer, Buffer]):
 Codec = ArrayArrayCodec | ArrayBytesCodec | BytesBytesCodec
 
 
+class CodecConfigDict(TypedDict):
+    """A dictionary representing a codec configuration."""
+
+    ...
+
+
+T = TypeVar("T", bound=CodecConfigDict)
+
+
+class CodecDict(TypedDict, Generic[T]):
+    """A generic dictionary representing a codec."""
+
+    name: str
+    configuration: NotRequired[T]
+
+
 class ArrayBytesCodecPartialDecodeMixin:
     """Mixin for array-to-bytes codecs that implement partial decoding."""
 

src/zarr/codecs/blosc.py

@@ -3,12 +3,12 @@
 from dataclasses import dataclass, replace
 from enum import Enum
 from functools import cached_property
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 import numcodecs
 from numcodecs.blosc import Blosc
 
-from zarr.abc.codec import BytesBytesCodec
+from zarr.abc.codec import BytesBytesCodec, CodecConfigDict, CodecDict
 from zarr.core.buffer.cpu import as_numpy_array_wrapper
 from zarr.core.common import JSON, parse_enum, parse_named_configuration, to_thread
 from zarr.registry import register_codec
@@ -54,6 +54,22 @@ class BloscCname(Enum):
     zlib = "zlib"
 
 
+class BloscCodecConfigDict(CodecConfigDict):
+    """A dictionary representing a Blosc codec configuration."""
+
+    typesize: int
+    cname: BloscCname
+    clevel: int
+    shuffle: BloscShuffle
+    blocksize: int
+
+
+class BloscCodecDict(CodecDict[BloscCodecConfigDict]):
+    """A dictionary representing a Blosc codec."""
+
+    ...
+
+
 # See https://zarr.readthedocs.io/en/stable/tutorial.html#configuring-blosc
 numcodecs.blosc.use_threads = False
 
@@ -118,12 +134,12 @@ def from_dict(cls, data: dict[str, JSON]) -> Self:
         _, configuration_parsed = parse_named_configuration(data, "blosc")
         return cls(**configuration_parsed)  # type: ignore[arg-type]
 
-    def to_dict(self) -> dict[str, JSON]:
+    def to_dict(self) -> BloscCodecDict:
         if self.typesize is None:
             raise ValueError("`typesize` needs to be set for serialization.")
         if self.shuffle is None:
             raise ValueError("`shuffle` needs to be set for serialization.")
-        return {
+        out_dict = {
             "name": "blosc",
             "configuration": {
                 "typesize": self.typesize,
@@ -134,6 +150,8 @@ def to_dict(self) -> dict[str, JSON]:
             },
         }
 
+        return cast(BloscCodecDict, out_dict)
+
     def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self:
         dtype = array_spec.dtype
         new_codec = self

src/zarr/codecs/bytes.py

@@ -7,7 +7,7 @@
 
 import numpy as np
 
-from zarr.abc.codec import ArrayBytesCodec
+from zarr.abc.codec import ArrayBytesCodec, CodecConfigDict, CodecDict
 from zarr.core.buffer import Buffer, NDArrayLike, NDBuffer
 from zarr.core.common import JSON, parse_enum, parse_named_configuration
 from zarr.registry import register_codec
@@ -30,6 +30,18 @@ class Endian(Enum):
 default_system_endian = Endian(sys.byteorder)
 
 
+class BytesCodecConfigDict(CodecConfigDict):
+    """A dictionary representing a bytes codec configuration."""
+
+    endian: Endian
+
+
+class BytesCodecDict(CodecDict[BytesCodecConfigDict]):
+    """A dictionary representing a bytes codec."""
+
+    ...
+
+
 @dataclass(frozen=True)
 class BytesCodec(ArrayBytesCodec):
     is_fixed_size = True
@@ -49,11 +61,12 @@ def from_dict(cls, data: dict[str, JSON]) -> Self:
         configuration_parsed = configuration_parsed or {}
         return cls(**configuration_parsed)  # type: ignore[arg-type]
 
-    def to_dict(self) -> dict[str, JSON]:
-        if self.endian is None:
-            return {"name": "bytes"}
-        else:
-            return {"name": "bytes", "configuration": {"endian": self.endian.value}}
+    def to_dict(self) -> BytesCodecDict:
+        out_dict: BytesCodecDict = {"name": "bytes"}
+        if self.endian is not None:
+            out_dict["configuration"] = {"endian": self.endian.value}
+
+        return out_dict
 
     def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self:
         if array_spec.dtype.itemsize == 0:

src/zarr/codecs/crc32c_.py

@@ -7,7 +7,7 @@
 import typing_extensions
 from crc32c import crc32c
 
-from zarr.abc.codec import BytesBytesCodec
+from zarr.abc.codec import BytesBytesCodec, CodecConfigDict, CodecDict
 from zarr.core.common import JSON, parse_named_configuration
 from zarr.registry import register_codec
 
@@ -18,6 +18,12 @@
     from zarr.core.buffer import Buffer
 
 
+class Crc32cCodecDict(CodecDict[CodecConfigDict]):
+    """A dictionary representing a CRC32C codec."""
+
+    ...
+
+
 @dataclass(frozen=True)
 class Crc32cCodec(BytesBytesCodec):
     is_fixed_size = True
@@ -27,8 +33,9 @@ def from_dict(cls, data: dict[str, JSON]) -> Self:
         parse_named_configuration(data, "crc32c", require_configuration=False)
         return cls()
 
-    def to_dict(self) -> dict[str, JSON]:
-        return {"name": "crc32c"}
+    def to_dict(self) -> Crc32cCodecDict:
+        out_dict = {"name": "crc32c"}
+        return cast(Crc32cCodecDict, out_dict)
 
     async def _decode_single(
         self,

src/zarr/codecs/gzip.py

@@ -1,11 +1,11 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 from numcodecs.gzip import GZip
 
-from zarr.abc.codec import BytesBytesCodec
+from zarr.abc.codec import BytesBytesCodec, CodecConfigDict, CodecDict
 from zarr.core.buffer.cpu import as_numpy_array_wrapper
 from zarr.core.common import JSON, parse_named_configuration, to_thread
 from zarr.registry import register_codec
@@ -17,6 +17,18 @@
     from zarr.core.buffer import Buffer
 
 
+class GzipCodecConfigDict(CodecConfigDict):
+    """A dictionary representing a gzip codec configuration."""
+
+    level: int
+
+
+class GzipCodecDict(CodecDict[GzipCodecConfigDict]):
+    """A dictionary representing a gzip codec."""
+
+    ...
+
+
 def parse_gzip_level(data: JSON) -> int:
     if not isinstance(data, (int)):
         raise TypeError(f"Expected int, got {type(data)}")
@@ -43,8 +55,9 @@ def from_dict(cls, data: dict[str, JSON]) -> Self:
         _, configuration_parsed = parse_named_configuration(data, "gzip")
         return cls(**configuration_parsed)  # type: ignore[arg-type]
 
-    def to_dict(self) -> dict[str, JSON]:
-        return {"name": "gzip", "configuration": {"level": self.level}}
+    def to_dict(self) -> GzipCodecDict:
+        out_dict = {"name": "gzip", "configuration": {"level": self.level}}
+        return cast(GzipCodecDict, out_dict)
 
     async def _decode_single(
         self,

src/zarr/codecs/sharding.py

@@ -15,6 +15,8 @@
     ArrayBytesCodecPartialDecodeMixin,
     ArrayBytesCodecPartialEncodeMixin,
     Codec,
+    CodecConfigDict,
+    CodecDict,
     CodecPipeline,
 )
 from zarr.abc.store import ByteGetter, ByteRangeRequest, ByteSetter
@@ -320,6 +322,21 @@ async def finalize(
         return await shard_builder.finalize(index_location, index_encoder)
 
 
+class ShardingCodecConfigDict(CodecConfigDict):
+    """A dictionary representing a sharding codec configuration."""
+
+    chunk_shape: list[int]  # TODO: Double check this
+    codecs: list[CodecDict[Any]]
+    index_codecs: list[CodecDict[Any]]
+    index_location: ShardingCodecIndexLocation
+
+
+class ShardingCodecDict(CodecDict[ShardingCodecConfigDict]):
+    """A dictionary representing a sharding codec."""
+
+    ...
+
+
 @dataclass(frozen=True)
 class ShardingCodec(
     ArrayBytesCodec, ArrayBytesCodecPartialDecodeMixin, ArrayBytesCodecPartialEncodeMixin
@@ -377,8 +394,8 @@ def from_dict(cls, data: dict[str, JSON]) -> Self:
     def codec_pipeline(self) -> CodecPipeline:
         return get_pipeline_class().from_codecs(self.codecs)
 
-    def to_dict(self) -> dict[str, JSON]:
-        return {
+    def to_dict(self) -> ShardingCodecDict:
+        out_dict = {
             "name": "sharding_indexed",
             "configuration": {
                 "chunk_shape": self.chunk_shape,
@@ -388,6 +405,8 @@ def to_dict(self) -> dict[str, JSON]:
             },
         }
 
+        return cast(ShardingCodecDict, out_dict)
+
     def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self:
         shard_spec = self._get_chunk_spec(array_spec)
         evolved_codecs = tuple(c.evolve_from_array_spec(array_spec=shard_spec) for c in self.codecs)

src/zarr/codecs/transpose.py

@@ -6,7 +6,7 @@
 
 import numpy as np
 
-from zarr.abc.codec import ArrayArrayCodec
+from zarr.abc.codec import ArrayArrayCodec, CodecConfigDict, CodecDict
 from zarr.core.array_spec import ArraySpec
 from zarr.core.common import JSON, ChunkCoordsLike, parse_named_configuration
 from zarr.registry import register_codec
@@ -26,6 +26,18 @@ def parse_transpose_order(data: JSON | Iterable[int]) -> tuple[int, ...]:
     return tuple(cast(Iterable[int], data))
 
 
+class TransposeCodecConfigDict(CodecConfigDict):
+    """A dictionary representing a transpose codec configuration."""
+
+    order: list[int]
+
+
+class TransposeCodecDict(CodecDict[TransposeCodecConfigDict]):
+    """A dictionary representing a transpose codec."""
+
+    ...
+
+
 @dataclass(frozen=True)
 class TransposeCodec(ArrayArrayCodec):
     is_fixed_size = True
@@ -42,8 +54,9 @@ def from_dict(cls, data: dict[str, JSON]) -> Self:
         _, configuration_parsed = parse_named_configuration(data, "transpose")
         return cls(**configuration_parsed)  # type: ignore[arg-type]
 
-    def to_dict(self) -> dict[str, JSON]:
-        return {"name": "transpose", "configuration": {"order": tuple(self.order)}}
+    def to_dict(self) -> TransposeCodecDict:
+        out_dict = {"name": "transpose", "configuration": {"order": tuple(self.order)}}
+        return cast(TransposeCodecDict, out_dict)
 
     def validate(self, shape: tuple[int, ...], dtype: np.dtype[Any], chunk_grid: ChunkGrid) -> None:
         if len(self.order) != len(shape):

src/zarr/codecs/zstd.py

@@ -3,11 +3,11 @@
 from dataclasses import dataclass
 from functools import cached_property
 from importlib.metadata import version
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 from numcodecs.zstd import Zstd
 
-from zarr.abc.codec import BytesBytesCodec
+from zarr.abc.codec import BytesBytesCodec, CodecConfigDict, CodecDict
 from zarr.core.buffer.cpu import as_numpy_array_wrapper
 from zarr.core.common import JSON, parse_named_configuration, to_thread
 from zarr.registry import register_codec
@@ -33,6 +33,19 @@ def parse_checksum(data: JSON) -> bool:
     raise TypeError(f"Expected bool. Got {type(data)}.")
 
 
+class ZstdCodecConfigDict(CodecConfigDict):
+    """A dictionary representing a zstd codec configuration."""
+
+    level: int
+    checksum: bool
+
+
+class ZstdCodecDict(CodecDict[ZstdCodecConfigDict]):
+    """A dictionary representing a zstd codec."""
+
+    ...
+
+
 @dataclass(frozen=True)
 class ZstdCodec(BytesBytesCodec):
     is_fixed_size = True
@@ -60,8 +73,12 @@ def from_dict(cls, data: dict[str, JSON]) -> Self:
         _, configuration_parsed = parse_named_configuration(data, "zstd")
         return cls(**configuration_parsed)  # type: ignore[arg-type]
 
-    def to_dict(self) -> dict[str, JSON]:
-        return {"name": "zstd", "configuration": {"level": self.level, "checksum": self.checksum}}
+    def to_dict(self) -> ZstdCodecDict:
+        out_dict = {
+            "name": "zstd",
+            "configuration": {"level": self.level, "checksum": self.checksum},
+        }
+        return cast(ZstdCodecDict, out_dict)
 
     @cached_property
     def _zstd_codec(self) -> Zstd: