Finish typing zarr.metadata

Author: dstansby

pyproject.toml

@@ -238,7 +238,6 @@ module = [
     "zarr.array",
     "zarr.common",
     "zarr.group",
-    "zarr.metadata"
 ]
 disallow_untyped_defs = false
 

src/zarr/chunk_grids.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import TYPE_CHECKING, Any, Dict
+from typing import TYPE_CHECKING, Dict
 from dataclasses import dataclass
 from zarr.abc.metadata import Metadata
 
@@ -18,13 +18,13 @@
 @dataclass(frozen=True)
 class ChunkGrid(Metadata):
     @classmethod
-    def from_dict(cls, data: Dict[str, JSON]) -> ChunkGrid:
+    def from_dict(cls, data: Dict[str, JSON] | ChunkGrid) -> ChunkGrid:
         if isinstance(data, ChunkGrid):
             return data
 
         name_parsed, _ = parse_named_configuration(data)
         if name_parsed == "regular":
-            return RegularChunkGrid.from_dict(data)
+            return RegularChunkGrid._from_dict(data)
         raise ValueError(f"Unknown chunk grid. Got {name_parsed}.")
 
 
@@ -38,7 +38,7 @@ def __init__(self, *, chunk_shape: ChunkCoordsLike) -> None:
         object.__setattr__(self, "chunk_shape", chunk_shape_parsed)
 
     @classmethod
-    def from_dict(cls, data: Dict[str, Any]) -> Self:
+    def _from_dict(cls, data: Dict[str, JSON]) -> Self:
         _, configuration_parsed = parse_named_configuration(data, "regular")
 
         return cls(**configuration_parsed)  # type: ignore[arg-type]

src/zarr/chunk_key_encodings.py

@@ -33,7 +33,7 @@ def __init__(self, *, separator: SeparatorLiteral) -> None:
         object.__setattr__(self, "separator", separator_parsed)
 
     @classmethod
-    def from_dict(cls, data: Dict[str, JSON]) -> ChunkKeyEncoding:
+    def from_dict(cls, data: Dict[str, JSON] | ChunkKeyEncoding) -> ChunkKeyEncoding:
         if isinstance(data, ChunkKeyEncoding):
             return data
 

src/zarr/common.py

@@ -19,6 +19,7 @@
     from typing import Awaitable, Callable, Iterator, Optional, Type
 
 import numpy as np
+import numpy.typing as npt
 
 ZARR_JSON = "zarr.json"
 ZARRAY_JSON = ".zarray"
@@ -152,7 +153,7 @@ def parse_named_configuration(
     return name_parsed, configuration_parsed
 
 
-def parse_shapelike(data: Any) -> tuple[int, ...]:
+def parse_shapelike(data: Iterable[int]) -> tuple[int, ...]:
     if not isinstance(data, Iterable):
         raise TypeError(f"Expected an iterable. Got {data} instead.")
     data_tuple = tuple(data)
@@ -166,7 +167,7 @@ def parse_shapelike(data: Any) -> tuple[int, ...]:
     return data_tuple
 
 
-def parse_dtype(data: Any) -> np.dtype[Any]:
+def parse_dtype(data: npt.DTypeLike) -> np.dtype[Any]:
     # todo: real validation
     return np.dtype(data)
 

src/zarr/metadata.py

@@ -17,6 +17,7 @@
 
 from zarr.abc.codec import Codec
 from zarr.abc.metadata import Metadata
+import numcodecs.abc
 
 from zarr.common import (
     JSON,
@@ -119,15 +120,15 @@ class ArrayMetadata(Metadata):
     def __init__(
         self,
         *,
-        shape,
-        data_type,
-        chunk_grid,
-        chunk_key_encoding,
-        fill_value,
-        codecs,
-        attributes,
-        dimension_names,
-    ):
+        shape: Iterable[int],
+        data_type: npt.DTypeLike,
+        chunk_grid: Dict[str, JSON] | ChunkGrid,
+        chunk_key_encoding: Dict[str, JSON] | ChunkKeyEncoding,
+        fill_value: Any,
+        codecs: Iterable[Union[Codec, JSON]],
+        attributes: None | Dict[str, JSON],
+        dimension_names: None | Iterable[str],
+    ) -> None:
         """
         Because the class is a frozen dataclass, we set attributes using object.__setattr__
         """
@@ -193,14 +194,14 @@ def get_chunk_spec(self, _chunk_coords: ChunkCoords, order: Literal["C", "F"]) -
         )
 
     def to_bytes(self) -> bytes:
-        def _json_convert(o):
+        def _json_convert(o: np.dtype[Any] | Enum | Codec) -> str | dict[str, Any]:
             if isinstance(o, np.dtype):
                 return str(o)
             if isinstance(o, Enum):
                 return o.name
             # this serializes numcodecs compressors
             # todo: implement to_dict for codecs
-            elif hasattr(o, "get_config"):
+            elif isinstance(o, numcodecs.abc.Codec):
                 return o.get_config()
             raise TypeError
 
@@ -290,7 +291,9 @@ def ndim(self) -> int:
         return len(self.shape)
 
     def to_bytes(self) -> bytes:
-        def _json_convert(o):
+        def _json_convert(
+            o: np.dtype[Any],
+        ) -> str | list[tuple[str, str] | tuple[str, str, tuple[int, ...]]]:
             if isinstance(o, np.dtype):
                 if o.fields is None:
                     return o.str
@@ -307,7 +310,7 @@ def from_dict(cls, data: Dict[str, Any]) -> ArrayV2Metadata:
         return cls(**data)
 
 
-def parse_dimension_names(data: Any) -> Tuple[str, ...] | None:
+def parse_dimension_names(data: None | Iterable[str]) -> Tuple[str, ...] | None:
     if data is None:
         return data
     if isinstance(data, Iterable) and all([isinstance(x, str) for x in data]):
@@ -317,12 +320,11 @@ def parse_dimension_names(data: Any) -> Tuple[str, ...] | None:
 
 
 # todo: real validation
-def parse_attributes(data: Any) -> Dict[str, JSON]:
+def parse_attributes(data: None | Dict[str, JSON]) -> Dict[str, JSON]:
     if data is None:
         return {}
 
-    data_json = cast(Dict[str, JSON], data)
-    return data_json
+    return data
 
 
 # todo: move to its own module and drop _v3 suffix