remove cache poisoing bug, and deploy type checker throughout the codebase

Author: d-v-b

examples/custom_dtype.py

@@ -29,7 +29,7 @@
     DTypeConfig_V2,
     DTypeJSON,
 )
-from zarr.core.type_check import check_type
+from zarr.core.type_check import guard_type
 
 # This is the int2 array data type
 int2_dtype_cls = type(np.dtype("int2"))
@@ -84,7 +84,7 @@ def _check_json_v2(cls, data: DTypeJSON) -> TypeGuard[DTypeConfig_V2[Literal["in
 
         See the Zarr docs for more information about the JSON encoding for data types.
         """
-        return check_type(data, DTypeConfig_V2[Literal["int2"], None]).success
+        return guard_type(data, DTypeConfig_V2[Literal["int2"], None])
 
     @classmethod
     def _check_json_v3(cls, data: DTypeJSON) -> TypeGuard[Literal["int2"]]:

src/zarr/abc/codec.py

@@ -1,14 +1,11 @@
 from __future__ import annotations
 
 from abc import abstractmethod
-from collections.abc import Mapping
-from typing import TYPE_CHECKING, Generic, TypeGuard, TypeVar
-
-from typing_extensions import ReadOnly, TypedDict
+from typing import TYPE_CHECKING, Generic, TypeVar
 
 from zarr.abc.metadata import Metadata
 from zarr.core.buffer import Buffer, NDBuffer
-from zarr.core.common import NamedConfig, concurrent_map
+from zarr.core.common import concurrent_map
 from zarr.core.config import config
 
 if TYPE_CHECKING:
@@ -37,27 +34,6 @@
 CodecInput = TypeVar("CodecInput", bound=NDBuffer | Buffer)
 CodecOutput = TypeVar("CodecOutput", bound=NDBuffer | Buffer)
 
-TName = TypeVar("TName", bound=str, covariant=True)
-
-
-class CodecJSON_V2(TypedDict, Generic[TName]):
-    """The JSON representation of a codec for Zarr V2"""
-
-    id: ReadOnly[TName]
-
-
-def _check_codecjson_v2(data: object) -> TypeGuard[CodecJSON_V2[str]]:
-    return isinstance(data, Mapping) and "id" in data and isinstance(data["id"], str)
-
-
-CodecJSON_V3 = str | NamedConfig[str, Mapping[str, object]]
-"""The JSON representation of a codec for Zarr V3."""
-
-# The widest type we will *accept* for a codec JSON
-# This covers v2 and v3
-CodecJSON = str | Mapping[str, object]
-"""The widest type of JSON-like input that could specify a codec."""
-
 
 class BaseCodec(Metadata, Generic[CodecInput, CodecOutput]):
     """Generic base class for codecs.

src/zarr/api/asynchronous.py

@@ -3,7 +3,7 @@
 import asyncio
 import dataclasses
 import warnings
-from typing import TYPE_CHECKING, Any, Literal, cast
+from typing import TYPE_CHECKING, Any, Literal
 
 import numpy as np
 import numpy.typing as npt
@@ -37,7 +37,7 @@
     GroupMetadata,
     create_hierarchy,
 )
-from zarr.core.metadata import ArrayMetadataDict, ArrayV2Metadata, ArrayV3Metadata
+from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata
 from zarr.errors import (
     GroupNotFoundError,
     NodeTypeValidationError,
@@ -352,13 +352,12 @@ async def open(
         try:
             metadata_dict = await get_array_metadata(store_path, zarr_format=zarr_format)
             # TODO: remove this cast when we fix typing for array metadata dicts
-            _metadata_dict = cast("ArrayMetadataDict", metadata_dict)
             # for v2, the above would already have raised an exception if not an array
-            zarr_format = _metadata_dict["zarr_format"]
-            is_v3_array = zarr_format == 3 and _metadata_dict.get("node_type") == "array"
+            zarr_format = metadata_dict["zarr_format"]
+            is_v3_array = zarr_format == 3 and metadata_dict.get("node_type") == "array"
             if is_v3_array or zarr_format == 2:
                 return AsyncArray(
-                    store_path=store_path, metadata=_metadata_dict, config=kwargs.get("config")
+                    store_path=store_path, metadata=metadata_dict, config=kwargs.get("config")
                 )
         except (AssertionError, FileNotFoundError, NodeTypeValidationError):
             pass

src/zarr/core/array.py

@@ -53,6 +53,8 @@
     ZARR_JSON,
     ZARRAY_JSON,
     ZATTRS_JSON,
+    ArrayMetadataJSON_V2,
+    ArrayMetadataJSON_V3,
     DimensionNames,
     MemoryOrder,
     ShapeLike,
@@ -102,11 +104,8 @@
 )
 from zarr.core.metadata import (
     ArrayMetadata,
-    ArrayMetadataDict,
     ArrayV2Metadata,
-    ArrayV2MetadataDict,
     ArrayV3Metadata,
-    ArrayV3MetadataDict,
     T_ArrayMetadata,
 )
 from zarr.core.metadata.v2 import (
@@ -115,9 +114,14 @@
     parse_compressor,
     parse_filters,
 )
-from zarr.core.metadata.v3 import parse_node_type_array
 from zarr.core.sync import sync
-from zarr.errors import MetadataValidationError, ZarrDeprecationWarning, ZarrUserWarning
+from zarr.core.type_check import check_type
+from zarr.errors import (
+    MetadataValidationError,
+    NodeTypeValidationError,
+    ZarrDeprecationWarning,
+    ZarrUserWarning,
+)
 from zarr.registry import (
     _parse_array_array_codec,
     _parse_array_bytes_codec,
@@ -176,12 +180,6 @@ def parse_array_metadata(data: Any) -> ArrayMetadata:
         zarr_format = data.get("zarr_format")
         if zarr_format == 3:
             meta_out = ArrayV3Metadata.from_dict(data)
-            if len(meta_out.storage_transformers) > 0:
-                msg = (
-                    f"Array metadata contains storage transformers: {meta_out.storage_transformers}."
-                    "Arrays with storage transformers are not supported in zarr-python at this time."
-                )
-                raise ValueError(msg)
             return meta_out
         elif zarr_format == 2:
             return ArrayV2Metadata.from_dict(data)
@@ -207,9 +205,27 @@ def create_codec_pipeline(metadata: ArrayMetadata, *, store: Store | None = None
     raise TypeError  # pragma: no cover
 
 
+@overload
+async def get_array_metadata(
+    store_path: StorePath, zarr_format: Literal[3]
+) -> ArrayMetadataJSON_V3: ...
+
+
+@overload
+async def get_array_metadata(
+    store_path: StorePath, zarr_format: Literal[2]
+) -> ArrayMetadataJSON_V2: ...
+
+
+@overload
+async def get_array_metadata(
+    store_path: StorePath, zarr_format: None
+) -> ArrayMetadataJSON_V3 | ArrayMetadataJSON_V2: ...
+
+
 async def get_array_metadata(
     store_path: StorePath, zarr_format: ZarrFormat | None = 3
-) -> dict[str, JSON]:
+) -> ArrayMetadataJSON_V3 | ArrayMetadataJSON_V2:
     if zarr_format == 2:
         zarray_bytes, zattrs_bytes = await gather(
             (store_path / ZARRAY_JSON).get(prototype=cpu_buffer_prototype),
@@ -241,19 +257,25 @@ async def get_array_metadata(
     else:
         raise MetadataValidationError("zarr_format", "2, 3, or None", zarr_format)
 
-    metadata_dict: dict[str, JSON]
+    metadata_dict: ArrayMetadataJSON_V2 | ArrayMetadataJSON_V3
     if zarr_format == 2:
         # V2 arrays are comprised of a .zarray and .zattrs objects
         assert zarray_bytes is not None
         metadata_dict = json.loads(zarray_bytes.to_bytes())
         zattrs_dict = json.loads(zattrs_bytes.to_bytes()) if zattrs_bytes is not None else {}
         metadata_dict["attributes"] = zattrs_dict
+        tycheck = check_type(metadata_dict, ArrayMetadataJSON_V2)
+        if not tycheck.success:
+            msg = "The .zarray object at {store_path} is not a valid Zarr array metadata object. "
+            raise NodeTypeValidationError("zarray", "Zarr array metadata object", metadata_dict)
     else:
         # V3 arrays are comprised of a zarr.json object
         assert zarr_json_bytes is not None
         metadata_dict = json.loads(zarr_json_bytes.to_bytes())
-
-        parse_node_type_array(metadata_dict.get("node_type"))
+        tycheck = check_type(metadata_dict, ArrayMetadataJSON_V3)
+        if not tycheck.success:
+            msg = "The zarr.json object at {store_path} is not a valid Zarr array metadata object. "
+            raise NodeTypeValidationError("zarr.json", "Zarr array metadata object", metadata_dict)
 
     return metadata_dict
 
@@ -292,22 +314,22 @@ class AsyncArray(Generic[T_ArrayMetadata]):
     @overload
     def __init__(
         self: AsyncArray[ArrayV2Metadata],
-        metadata: ArrayV2Metadata | ArrayV2MetadataDict,
+        metadata: ArrayV2Metadata | ArrayMetadataJSON_V2,
         store_path: StorePath,
         config: ArrayConfigLike | None = None,
     ) -> None: ...
 
     @overload
     def __init__(
         self: AsyncArray[ArrayV3Metadata],
-        metadata: ArrayV3Metadata | ArrayV3MetadataDict,
+        metadata: ArrayV3Metadata | ArrayMetadataJSON_V3,
         store_path: StorePath,
         config: ArrayConfigLike | None = None,
     ) -> None: ...
 
     def __init__(
         self,
-        metadata: ArrayMetadata | ArrayMetadataDict,
+        metadata: ArrayMetadata | ArrayMetadataJSON_V2 | ArrayMetadataJSON_V3,
         store_path: StorePath,
         config: ArrayConfigLike | None = None,
     ) -> None:
@@ -959,9 +981,7 @@ async def open(
         """
         store_path = await make_store_path(store)
         metadata_dict = await get_array_metadata(store_path, zarr_format=zarr_format)
-        # TODO: remove this cast when we have better type hints
-        _metadata_dict = cast("ArrayV3MetadataDict", metadata_dict)
-        return cls(store_path=store_path, metadata=_metadata_dict)
+        return cls(store_path=store_path, metadata=metadata_dict)
 
     @property
     def store(self) -> Store:

src/zarr/core/common.py

@@ -14,6 +14,7 @@
     Final,
     Generic,
     Literal,
+    NotRequired,
     TypedDict,
     TypeVar,
     cast,
@@ -47,11 +48,11 @@
 ANY_ACCESS_MODE: Final = "r", "r+", "a", "w", "w-"
 DimensionNames = Iterable[str | None] | None
 
-TName = TypeVar("TName", bound=str)
+TName = TypeVar("TName", bound=str, covariant=True)
 TConfig = TypeVar("TConfig", bound=Mapping[str, object])
 
 
-class NamedConfig(TypedDict, Generic[TName, TConfig]):
+class NamedRequiredConfig(TypedDict, Generic[TName, TConfig]):
     """
     A typed dictionary representing an object with a name and configuration, where the configuration
     is a mapping of string keys to values, e.g. another typed dictionary or a JSON object.
@@ -67,6 +68,104 @@ class NamedConfig(TypedDict, Generic[TName, TConfig]):
     """The configuration of the object."""
 
 
+class NamedConfig(TypedDict, Generic[TName, TConfig]):
+    """
+    A typed dictionary representing an object with a name and configuration, where the configuration
+    is a mapping of string keys to values, e.g. another typed dictionary or a JSON object.
+
+    The configuration key is not required.
+
+    This class is generic with two type parameters: the type of the name (``TName``) and the type of
+    the configuration (``TConfig``).
+    """
+
+    name: ReadOnly[TName]
+    """The name of the object."""
+
+    configuration: ReadOnly[NotRequired[TConfig]]
+    """The configuration of the object."""
+
+
+class ArrayMetadataJSON_V2(TypedDict):
+    """
+    A typed dictionary model for Zarr V2 array metadata.
+    """
+
+    zarr_format: Literal[2]
+    dtype: str | StructuredName_V2
+    shape: Sequence[int]
+    chunks: Sequence[int]
+    dimension_separator: NotRequired[Literal[".", "/"]]
+    fill_value: Any
+    filters: Sequence[CodecJSON_V2[str]] | None
+    order: Literal["C", "F"]
+    compressor: CodecJSON_V2[str] | None
+    attributes: NotRequired[Mapping[str, JSON]]
+
+
+class GroupMetadataJSON_V2(TypedDict):
+    """
+    A typed dictionary model for Zarr V2 group metadata.
+    """
+
+    zarr_format: Literal[2]
+    attributes: NotRequired[Mapping[str, JSON]]
+
+
+class ArrayMetadataJSON_V3(TypedDict):
+    """
+    A typed dictionary model for Zarr V3 array metadata.
+    """
+
+    zarr_format: Literal[3]
+    node_type: Literal["array"]
+    data_type: str | NamedConfig[str, Mapping[str, object]]
+    shape: Sequence[int]
+    chunk_grid: NamedConfig[str, Mapping[str, object]]
+    chunk_key_encoding: NamedConfig[str, Mapping[str, object]]
+    fill_value: object
+    codecs: Sequence[str | NamedConfig[str, Mapping[str, object]]]
+    attributes: NotRequired[Mapping[str, object]]
+    storage_transformers: NotRequired[Sequence[NamedConfig[str, Mapping[str, object]]]]
+    dimension_names: NotRequired[Sequence[str | None]]
+
+
+class GroupMetadataJSON_V3(TypedDict):
+    """
+    A typed dictionary model for Zarr V3 group metadata.
+    """
+
+    zarr_format: Literal[3]
+    node_type: Literal["group"]
+    attributes: NotRequired[Mapping[str, JSON]]
+
+
+class CodecJSON_V2(TypedDict, Generic[TName]):
+    """The JSON representation of a codec for Zarr V2"""
+
+    id: ReadOnly[TName]
+
+
+CodecJSON_V3 = str | NamedConfig[str, Mapping[str, object]]
+"""The JSON representation of a codec for Zarr V3."""
+
+# The widest type we will *accept* for a codec JSON
+# This covers v2 and v3
+CodecJSON = str | Mapping[str, object]
+"""The widest type of JSON-like input that could specify a codec."""
+
+
+# By comparison, The JSON representation of a dtype in zarr v3 is much simpler.
+# It's either a string, or a structured dict
+DTypeSpec_V3 = str | NamedConfig[str, Mapping[str, object]]
+
+#  This is the JSON representation of a structured dtype in zarr v2
+StructuredName_V2 = Sequence["str | StructuredName_V2"]
+
+# This models the type of the name a dtype might have in zarr v2 array metadata
+DTypeName_V2 = StructuredName_V2 | str
+
+
 def product(tup: tuple[int, ...]) -> int:
     return functools.reduce(operator.mul, tup, 1)