fix: accept extra fields in array metadata (#3530)

* add support for reading and preserving unknown, permitted fields in array metadata.

* sort in error message

* sort keys in the right place

* changelog

* rename changelog

Author: d-v-b

changes/3530.bugfix.md

@@ -0,0 +1,2 @@
+Fixed a bug that prevented Zarr Python from opening Zarr V3 array metadata documents that contained
+ extra keys with permissible values (dicts with a `"must_understand"` key set to `"false"`).

src/zarr/core/array.py

@@ -103,10 +103,10 @@
 from zarr.core.metadata import (
     ArrayMetadata,
     ArrayMetadataDict,
+    ArrayMetadataJSON_V3,
     ArrayV2Metadata,
     ArrayV2MetadataDict,
     ArrayV3Metadata,
-    ArrayV3MetadataDict,
     T_ArrayMetadata,
 )
 from zarr.core.metadata.io import save_metadata
@@ -319,7 +319,7 @@ def __init__(
     @overload
     def __init__(
         self: AsyncArray[ArrayV3Metadata],
-        metadata: ArrayV3Metadata | ArrayV3MetadataDict,
+        metadata: ArrayV3Metadata | ArrayMetadataJSON_V3,
         store_path: StorePath,
         config: ArrayConfigLike | None = None,
     ) -> None: ...
@@ -1004,7 +1004,7 @@ async def example():
         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)
+        _metadata_dict = cast("ArrayMetadataJSON_V3", metadata_dict)
         return cls(store_path=store_path, metadata=_metadata_dict)
 
     @property

src/zarr/core/chunk_grids.py

@@ -15,6 +15,7 @@
 from zarr.abc.metadata import Metadata
 from zarr.core.common import (
     JSON,
+    NamedConfig,
     ShapeLike,
     ceildiv,
     parse_named_configuration,
@@ -152,7 +153,7 @@ def normalize_chunks(chunks: Any, shape: tuple[int, ...], typesize: int) -> tupl
 @dataclass(frozen=True)
 class ChunkGrid(Metadata):
     @classmethod
-    def from_dict(cls, data: dict[str, JSON] | ChunkGrid) -> ChunkGrid:
+    def from_dict(cls, data: dict[str, JSON] | ChunkGrid | NamedConfig[str, Any]) -> ChunkGrid:
         if isinstance(data, ChunkGrid):
             return data
 
@@ -180,7 +181,7 @@ def __init__(self, *, chunk_shape: ShapeLike) -> None:
         object.__setattr__(self, "chunk_shape", chunk_shape_parsed)
 
     @classmethod
-    def _from_dict(cls, data: dict[str, JSON]) -> Self:
+    def _from_dict(cls, data: dict[str, JSON] | NamedConfig[str, Any]) -> Self:
         _, configuration_parsed = parse_named_configuration(data, "regular")
 
         return cls(**configuration_parsed)  # type: ignore[arg-type]

src/zarr/core/chunk_key_encodings.py

@@ -2,14 +2,15 @@
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, ClassVar, Literal, TypeAlias, TypedDict, cast
+from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeAlias, TypedDict, cast
 
 if TYPE_CHECKING:
     from typing import NotRequired, Self
 
 from zarr.abc.metadata import Metadata
 from zarr.core.common import (
     JSON,
+    NamedConfig,
     parse_named_configuration,
 )
 from zarr.registry import get_chunk_key_encoding_class, register_chunk_key_encoding
@@ -61,7 +62,9 @@ def encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str:
         """
 
 
-ChunkKeyEncodingLike: TypeAlias = dict[str, JSON] | ChunkKeyEncodingParams | ChunkKeyEncoding
+ChunkKeyEncodingLike: TypeAlias = (
+    dict[str, JSON] | ChunkKeyEncodingParams | ChunkKeyEncoding | NamedConfig[str, Any]
+)
 
 
 @dataclass(frozen=True)
@@ -108,7 +111,7 @@ def parse_chunk_key_encoding(data: ChunkKeyEncodingLike) -> ChunkKeyEncoding:
 
     # handle ChunkKeyEncodingParams
     if "name" in data and "separator" in data:
-        data = {"name": data["name"], "configuration": {"separator": data["separator"]}}
+        data = {"name": data["name"], "configuration": {"separator": data["separator"]}}  # type: ignore[typeddict-item]
 
     # Now must be a named config
     data = cast("dict[str, JSON]", data)

src/zarr/core/common.py

@@ -14,6 +14,7 @@
     Final,
     Generic,
     Literal,
+    NotRequired,
     TypedDict,
     TypeVar,
     cast,
@@ -63,8 +64,8 @@ class NamedConfig(TypedDict, Generic[TName, TConfig]):
     name: ReadOnly[TName]
     """The name of the object."""
 
-    configuration: ReadOnly[TConfig]
-    """The configuration of the object."""
+    configuration: NotRequired[ReadOnly[TConfig]]
+    """The configuration of the object. Not required."""
 
 
 def product(tup: tuple[int, ...]) -> int:
@@ -134,18 +135,24 @@ def parse_configuration(data: JSON) -> JSON:
 
 @overload
 def parse_named_configuration(
-    data: JSON, expected_name: str | None = None
+    data: JSON | NamedConfig[str, Any], expected_name: str | None = None
 ) -> tuple[str, dict[str, JSON]]: ...
 
 
 @overload
 def parse_named_configuration(
-    data: JSON, expected_name: str | None = None, *, require_configuration: bool = True
+    data: JSON | NamedConfig[str, Any],
+    expected_name: str | None = None,
+    *,
+    require_configuration: bool = True,
 ) -> tuple[str, dict[str, JSON] | None]: ...
 
 
 def parse_named_configuration(
-    data: JSON, expected_name: str | None = None, *, require_configuration: bool = True
+    data: JSON | NamedConfig[str, Any],
+    expected_name: str | None = None,
+    *,
+    require_configuration: bool = True,
 ) -> tuple[str, JSON | None]:
     if not isinstance(data, dict):
         raise TypeError(f"Expected dict, got {type(data)}")

src/zarr/core/metadata/__init__.py

@@ -1,17 +1,17 @@
 from typing import TypeAlias, TypeVar
 
 from .v2 import ArrayV2Metadata, ArrayV2MetadataDict
-from .v3 import ArrayV3Metadata, ArrayV3MetadataDict
+from .v3 import ArrayMetadataJSON_V3, ArrayV3Metadata
 
 ArrayMetadata: TypeAlias = ArrayV2Metadata | ArrayV3Metadata
-ArrayMetadataDict: TypeAlias = ArrayV2MetadataDict | ArrayV3MetadataDict
+ArrayMetadataDict: TypeAlias = ArrayV2MetadataDict | ArrayMetadataJSON_V3
 T_ArrayMetadata = TypeVar("T_ArrayMetadata", ArrayV2Metadata, ArrayV3Metadata)
 
 __all__ = [
     "ArrayMetadata",
     "ArrayMetadataDict",
+    "ArrayMetadataJSON_V3",
     "ArrayV2Metadata",
     "ArrayV2MetadataDict",
     "ArrayV3Metadata",
-    "ArrayV3MetadataDict",
 ]

src/zarr/core/metadata/v3.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, TypedDict
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, NotRequired, TypedDict, TypeGuard, cast
 
 from zarr.abc.metadata import Metadata
 from zarr.core.buffer.core import default_buffer_prototype
@@ -33,6 +34,7 @@
     JSON,
     ZARR_JSON,
     DimensionNames,
+    NamedConfig,
     parse_named_configuration,
     parse_shapelike,
 )
@@ -136,13 +138,61 @@ def parse_storage_transformers(data: object) -> tuple[dict[str, JSON], ...]:
     )
 
 
-class ArrayV3MetadataDict(TypedDict):
+class AllowedExtraField(TypedDict):
+    """
+    This class models allowed extra fields in array metadata.
+    They are ignored by Zarr Python.
+    """
+
+    must_understand: Literal[False]
+
+
+def check_allowed_extra_field(data: object) -> TypeGuard[AllowedExtraField]:
+    """
+    Check if the extra field is allowed according to the Zarr v3 spec. The object
+    must be a mapping with a "must_understand" key set to `False`.
+    """
+    return isinstance(data, Mapping) and data.get("must_understand") is False
+
+
+def parse_extra_fields(
+    data: Mapping[str, AllowedExtraField] | None,
+) -> dict[str, AllowedExtraField]:
+    if data is None:
+        return {}
+    else:
+        conflict_keys = ARRAY_METADATA_KEYS & set(data.keys())
+        if len(conflict_keys) > 0:
+            msg = (
+                "Invalid extra fields. "
+                "The following keys: "
+                f"{sorted(conflict_keys)} "
+                "are invalid because they collide with keys reserved for use by the "
+                "array metadata document."
+            )
+            raise ValueError(msg)
+        return dict(data)
+
+
+class ArrayMetadataJSON_V3(TypedDict):
     """
     A typed dictionary model for zarr v3 metadata.
     """
 
     zarr_format: Literal[3]
-    attributes: dict[str, JSON]
+    node_type: Literal["array"]
+    data_type: str | NamedConfig[str, Mapping[str, object]]
+    shape: tuple[int, ...]
+    chunk_grid: NamedConfig[str, Mapping[str, object]]
+    chunk_key_encoding: NamedConfig[str, Mapping[str, object]]
+    fill_value: object
+    codecs: tuple[str | NamedConfig[str, Mapping[str, object]], ...]
+    attributes: NotRequired[Mapping[str, JSON]]
+    storage_transformers: NotRequired[tuple[NamedConfig[str, Mapping[str, object]], ...]]
+    dimension_names: NotRequired[tuple[str | None]]
+
+
+ARRAY_METADATA_KEYS = set(ArrayMetadataJSON_V3.__annotations__.keys())
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -158,19 +208,21 @@ class ArrayV3Metadata(Metadata):
     zarr_format: Literal[3] = field(default=3, init=False)
     node_type: Literal["array"] = field(default="array", init=False)
     storage_transformers: tuple[dict[str, JSON], ...]
+    extra_fields: dict[str, AllowedExtraField]
 
     def __init__(
         self,
         *,
         shape: Iterable[int],
         data_type: ZDType[TBaseDType, TBaseScalar],
-        chunk_grid: dict[str, JSON] | ChunkGrid,
+        chunk_grid: dict[str, JSON] | ChunkGrid | NamedConfig[str, Any],
         chunk_key_encoding: ChunkKeyEncodingLike,
         fill_value: object,
-        codecs: Iterable[Codec | dict[str, JSON]],
+        codecs: Iterable[Codec | dict[str, JSON] | NamedConfig[str, Any] | str],
         attributes: dict[str, JSON] | None,
         dimension_names: DimensionNames,
         storage_transformers: Iterable[dict[str, JSON]] | None = None,
+        extra_fields: Mapping[str, AllowedExtraField] | None = None,
     ) -> None:
         """
         Because the class is a frozen dataclass, we set attributes using object.__setattr__
@@ -185,7 +237,7 @@ def __init__(
         attributes_parsed = parse_attributes(attributes)
         codecs_parsed_partial = parse_codecs(codecs)
         storage_transformers_parsed = parse_storage_transformers(storage_transformers)
-
+        extra_fields_parsed = parse_extra_fields(extra_fields)
         array_spec = ArraySpec(
             shape=shape_parsed,
             dtype=data_type,
@@ -205,6 +257,7 @@ def __init__(
         object.__setattr__(self, "fill_value", fill_value_parsed)
         object.__setattr__(self, "attributes", attributes_parsed)
         object.__setattr__(self, "storage_transformers", storage_transformers_parsed)
+        object.__setattr__(self, "extra_fields", extra_fields_parsed)
 
         self._validate_metadata()
 
@@ -323,16 +376,45 @@ def from_dict(cls, data: dict[str, JSON]) -> Self:
         except ValueError as e:
             raise TypeError(f"Invalid fill_value: {fill!r}") from e
 
-        # dimension_names key is optional, normalize missing to `None`
-        _data["dimension_names"] = _data.pop("dimension_names", None)
-
-        # attributes key is optional, normalize missing to `None`
-        _data["attributes"] = _data.pop("attributes", None)
-
-        return cls(**_data, fill_value=fill_value_parsed, data_type=data_type)  # type: ignore[arg-type]
+        # check if there are extra keys
+        extra_keys = set(_data.keys()) - ARRAY_METADATA_KEYS
+        allowed_extra_fields: dict[str, AllowedExtraField] = {}
+        invalid_extra_fields = {}
+        for key in extra_keys:
+            val = _data[key]
+            if check_allowed_extra_field(val):
+                allowed_extra_fields[key] = val
+            else:
+                invalid_extra_fields[key] = val
+        if len(invalid_extra_fields) > 0:
+            msg = (
+                "Got a Zarr V3 metadata document with the following disallowed extra fields:"
+                f"{sorted(invalid_extra_fields.keys())}."
+                'Extra fields are not allowed unless they are a dict with a "must_understand" key'
+                "which is assigned the value `False`."
+            )
+            raise MetadataValidationError(msg)
+        # TODO: replace this with a real type check!
+        _data_typed = cast(ArrayMetadataJSON_V3, _data)
+
+        return cls(
+            shape=_data_typed["shape"],
+            chunk_grid=_data_typed["chunk_grid"],
+            chunk_key_encoding=_data_typed["chunk_key_encoding"],
+            codecs=_data_typed["codecs"],
+            attributes=_data_typed.get("attributes", {}),  # type: ignore[arg-type]
+            dimension_names=_data_typed.get("dimension_names", None),
+            fill_value=fill_value_parsed,
+            data_type=data_type,
+            extra_fields=allowed_extra_fields,
+            storage_transformers=_data_typed.get("storage_transformers", ()),  # type: ignore[arg-type]
+        )
 
     def to_dict(self) -> dict[str, JSON]:
         out_dict = super().to_dict()
+        extra_fields = out_dict.pop("extra_fields")
+        out_dict = out_dict | extra_fields  # type: ignore[operator]
+
         out_dict["fill_value"] = self.data_type.to_json_scalar(
             self.fill_value, zarr_format=self.zarr_format
         )
@@ -351,7 +433,6 @@ def to_dict(self) -> dict[str, JSON]:
         dtype_meta = out_dict["data_type"]
         if isinstance(dtype_meta, ZDType):
             out_dict["data_type"] = dtype_meta.to_json(zarr_format=3)  # type: ignore[unreachable]
-
         return out_dict
 
     def update_shape(self, shape: tuple[int, ...]) -> Self: