fix(group): accept must_understand=false extra fields in v3 group metadata

changes/4451.bugfix.md

@@ -0,0 +1 @@
+Zarr format 3 group metadata now accepts extra fields that are objects with `"must_understand": false`, as the specification requires, and writes them back unchanged; previously opening such a group raised `TypeError`. Extra fields of any other shape raise `MetadataValidationError`, both when reading a metadata document and when constructing `GroupMetadata` or `ArrayV3Metadata` directly. `GroupMetadata` exposes these fields as `extra_fields`, like `ArrayV3Metadata`.

src/zarr/core/group.py

@@ -5,9 +5,9 @@
 import unicodedata
 import warnings
 from collections import defaultdict
-from dataclasses import asdict, dataclass, field, fields, replace
+from dataclasses import asdict, dataclass, field, replace
 from itertools import accumulate
-from typing import TYPE_CHECKING, Literal, assert_never, cast, overload
+from typing import TYPE_CHECKING, Final, Literal, assert_never, cast, overload
 
 import numpy as np
 
@@ -49,6 +49,7 @@
 from zarr.core.json_parse import parse_field
 from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata
 from zarr.core.metadata.io import save_metadata
+from zarr.core.metadata.v3 import AllowedExtraField, parse_extra_fields
 from zarr.core.sync import SyncMixin, sync
 from zarr.errors import (
     ArrayNotFoundError,
@@ -311,9 +312,9 @@ def flattened_metadata(self) -> dict[str, ArrayV2Metadata | ArrayV3Metadata | Gr
                 "group-1": GroupMetadata(),
             }
         )
-        # {'group-0': GroupMetadata(attributes={}, zarr_format=3, consolidated_metadata=None, node_type='group'),
-        #  'group-0/group-0-0': GroupMetadata(attributes={}, zarr_format=3, consolidated_metadata=None, node_type='group'),
-        #  'group-1': GroupMetadata(attributes={}, zarr_format=3, consolidated_metadata=None, node_type='group')}
+        # {'group-0': GroupMetadata(attributes={}, zarr_format=3, consolidated_metadata=None, node_type='group', extra_fields={}),
+        #  'group-0/group-0-0': GroupMetadata(attributes={}, zarr_format=3, consolidated_metadata=None, node_type='group', extra_fields={}),
+        #  'group-1': GroupMetadata(attributes={}, zarr_format=3, consolidated_metadata=None, node_type='group', extra_fields={})}
         ```
         """
         metadata = {}
@@ -345,6 +346,18 @@ def flatten(
         return metadata
 
 
+GROUP_METADATA_KEYS: Final[set[str]] = {
+    "zarr_format",
+    "node_type",
+    "attributes",
+    "consolidated_metadata",
+}
+"""
+The names of the fields of the group metadata document. `consolidated_metadata` is not in the
+Zarr V3 spec; Zarr-Python writes it.
+"""
+
+
 @dataclass(frozen=True)
 class GroupMetadata(Metadata):
     """
@@ -355,6 +368,7 @@ class GroupMetadata(Metadata):
     zarr_format: ZarrFormat = 3
     consolidated_metadata: ConsolidatedMetadata | None = None
     node_type: Literal["group"] = field(default="group", init=False)
+    extra_fields: dict[str, AllowedExtraField] = field(default_factory=dict)
 
     def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]:
         indent = config.get("json_indent")
@@ -406,13 +420,22 @@ def __init__(
         attributes: dict[str, Any] | None = None,
         zarr_format: ZarrFormat = 3,
         consolidated_metadata: ConsolidatedMetadata | None = None,
+        extra_fields: Mapping[str, AllowedExtraField] | None = None,
     ) -> None:
         attributes_parsed = parse_attributes(attributes)
         zarr_format_parsed = parse_zarr_format(zarr_format)
+        if zarr_format_parsed == 2 and extra_fields:
+            raise ValueError(
+                "Invalid extra fields. Zarr format 2 group metadata does not support extra fields."
+            )
+        extra_fields_parsed = parse_extra_fields(
+            extra_fields, reserved_keys=GROUP_METADATA_KEYS, node_type="group"
+        )
 
         object.__setattr__(self, "attributes", attributes_parsed)
         object.__setattr__(self, "zarr_format", zarr_format_parsed)
         object.__setattr__(self, "consolidated_metadata", consolidated_metadata)
+        object.__setattr__(self, "extra_fields", extra_fields_parsed)
 
     @classmethod
     def from_dict(cls, data: dict[str, Any]) -> GroupMetadata:
@@ -427,17 +450,22 @@ def from_dict(cls, data: dict[str, Any]) -> GroupMetadata:
             data["consolidated_metadata"] = ConsolidatedMetadata.from_dict(consolidated_metadata)
 
         zarr_format = data.get("zarr_format")
-        if zarr_format == 2 or zarr_format is None:
-            # zarr v2 allowed arbitrary keys here.
-            # We don't want the GroupMetadata constructor to fail just because someone put an
-            # extra key in the metadata.
-            expected = {x.name for x in fields(cls)}
-            data = {k: v for k, v in data.items() if k in expected}
-
-        return cls(**data)
+        if zarr_format == 3:
+            # Zarr v3 allows an extra key only if it is an object with "must_understand": false.
+            extra_fields = parse_extra_fields(
+                {k: v for k, v in data.items() if k not in GROUP_METADATA_KEYS},
+                reserved_keys=GROUP_METADATA_KEYS,
+                node_type="group",
+            )
+        else:
+            # Zarr v2 allowed arbitrary keys; they are dropped rather than kept.
+            extra_fields = {}
+        data = {k: v for k, v in data.items() if k in GROUP_METADATA_KEYS}
+        return cls(**data, extra_fields=extra_fields)
 
     def to_dict(self) -> dict[str, Any]:
         result = asdict(replace(self, consolidated_metadata=None))
+        result.update(result.pop("extra_fields"))
         if self.consolidated_metadata is not None:
             result["consolidated_metadata"] = self.consolidated_metadata.to_dict()
         else:

src/zarr/core/metadata/v3.py

@@ -2,6 +2,7 @@
 
 import json
 from collections.abc import Iterable, Mapping, Sequence
+from collections.abc import Set as AbstractSet
 from dataclasses import dataclass, field, replace
 from typing import TYPE_CHECKING, Any, Final, Literal, NotRequired, TypeGuard, cast
 
@@ -24,6 +25,7 @@
     DimensionNamesLike,
     NamedConfig,
     NamedRequiredConfig,
+    NodeType,
     compress_rle,
     expand_rle,
     parse_named_configuration,
@@ -155,8 +157,8 @@ def parse_storage_transformers(data: object) -> tuple[dict[str, JSON], ...]:
 
 class AllowedExtraField(TypedDict, extra_items=JSON):  # type: ignore[call-arg]
     """
-    This class models allowed extra fields in array metadata.
-    They must have ``must_understand`` set to ``False``, and may contain
+    This class models allowed extra fields in array or group metadata.
+    They must have `must_understand` set to `False`, and may contain
     arbitrary additional JSON data.
     """
 
@@ -172,22 +174,45 @@ def check_allowed_extra_field(data: object) -> TypeGuard[AllowedExtraField]:
 
 
 def parse_extra_fields(
-    data: Mapping[str, AllowedExtraField] | None,
+    data: Mapping[str, object] | None,
+    *,
+    reserved_keys: AbstractSet[str],
+    node_type: NodeType,
 ) -> dict[str, AllowedExtraField]:
+    """
+    Check the extra fields of a Zarr V3 metadata document.
+
+    Raises `ValueError` if a key collides with a key in `reserved_keys`, and
+    `MetadataValidationError` if a value is not an object with `"must_understand": false`.
+    """
     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)
+    conflict_keys = reserved_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 "
+            f"{node_type} metadata document."
+        )
+        raise ValueError(msg)
+    allowed_extra_fields: dict[str, AllowedExtraField] = {}
+    invalid_extra_fields: list[str] = []
+    for key, val in data.items():
+        if check_allowed_extra_field(val):
+            allowed_extra_fields[key] = val
+        else:
+            invalid_extra_fields.append(key)
+    if len(invalid_extra_fields) > 0:
+        msg = (
+            f"Got Zarr V3 {node_type} metadata with the following disallowed extra "
+            f"fields: {sorted(invalid_extra_fields)}. "
+            'Extra fields are not allowed unless they are an object with a "must_understand" '
+            "key which is assigned the value `false`."
+        )
+        raise MetadataValidationError(msg)
+    return allowed_extra_fields
 
 
 # JSON type for a single dimension's rectilinear spec:
@@ -518,7 +543,9 @@ 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)
+        extra_fields_parsed = parse_extra_fields(
+            extra_fields, reserved_keys=ARRAY_METADATA_KEYS, node_type="array"
+        )
         array_spec = ArraySpec(
             shape=shape_parsed,
             dtype=data_type,
@@ -652,23 +679,11 @@ def from_dict(cls, data: dict[str, JSON]) -> Self:
             raise TypeError(f"Invalid fill_value: {fill!r}") from e
 
         # 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)
+        allowed_extra_fields = parse_extra_fields(
+            {k: v for k, v in _data.items() if k not in ARRAY_METADATA_KEYS},
+            reserved_keys=ARRAY_METADATA_KEYS,
+            node_type="array",
+        )
         # TODO: replace this with a real type check!
         _data_typed = cast(ArrayMetadataJSON_V3, _data)
 

src/zarr/core/sync_group.py

@@ -103,7 +103,7 @@ def create_hierarchy(
     nodes = {'a': GroupMetadata(attributes={'name': 'leaf'})}
     nodes_created = dict(create_hierarchy(store=store, nodes=nodes))
     print(nodes)
-    # {'a': GroupMetadata(attributes={'name': 'leaf'}, zarr_format=3, consolidated_metadata=None, node_type='group')}
+    # {'a': GroupMetadata(attributes={'name': 'leaf'}, zarr_format=3, consolidated_metadata=None, node_type='group', extra_fields={})}
     ```
     """
     coro = create_hierarchy_async(store=store, nodes=nodes, overwrite=overwrite)

tests/test_group.py

@@ -28,6 +28,7 @@
 from zarr.core.dtype.common import unpack_dtype_json
 from zarr.core.dtype.npy.int import UInt8
 from zarr.core.group import (
+    GROUP_METADATA_KEYS,
     ConsolidatedMetadata,
     GroupMetadata,
     ImplicitGroupMarker,
@@ -1718,12 +1719,75 @@ def test_from_dict_extra_fields(self):
         data = {
             "attributes": {"key": "value"},
             "_nczarr_superblock": {"version": "2.0.0"},
+            # a key named like the `extra_fields` parameter is dropped, not passed to __init__
+            "extra_fields": {"my_extension": {"must_understand": False}},
             "zarr_format": 2,
         }
         result = GroupMetadata.from_dict(data)
         expected = GroupMetadata(attributes={"key": "value"}, zarr_format=2)
         assert result == expected
 
+    @pytest.mark.parametrize(
+        "extra_fields",
+        [
+            {},
+            {"my_extension": {"must_understand": False, "version": 1}},
+            {"a": {"must_understand": False}, "b": {"must_understand": False, "data": [1, 2]}},
+            {"extra_fields": {"must_understand": False}},
+        ],
+    )
+    def test_v3_extra_fields_round_trip(self, extra_fields: dict[str, Any]) -> None:
+        # https://github.com/zarr-developers/zarr-python/issues/3523
+        data = {"zarr_format": 3, "node_type": "group", "attributes": {}, **extra_fields}
+        expected = GroupMetadata(extra_fields=extra_fields)
+        result = GroupMetadata.from_dict(data)
+        assert result == expected
+        assert result.extra_fields == extra_fields
+        assert result.to_dict() == data
+
+    @pytest.mark.parametrize("value", [{"must_understand": True}, {}, "not an object", 1])
+    def test_from_dict_v3_disallowed_extra_fields(self, value: object) -> None:
+        data = {"zarr_format": 3, "node_type": "group", "my_extension": value}
+        with pytest.raises(MetadataValidationError, match="my_extension"):
+            GroupMetadata.from_dict(data)
+
+    @pytest.mark.parametrize("value", [{"must_understand": True}, {}, "not an object", 1])
+    def test_init_disallowed_extra_fields(self, value: object) -> None:
+        with pytest.raises(MetadataValidationError, match="my_extension"):
+            GroupMetadata(extra_fields={"my_extension": value})  # type: ignore[dict-item]
+
+    @pytest.mark.parametrize("key", sorted(GROUP_METADATA_KEYS))
+    def test_init_extra_fields_collision(self, key: str) -> None:
+        extra_fields: dict[str, Any] = {key: {"must_understand": False}}
+        with pytest.raises(ValueError, match="collide with keys reserved"):
+            GroupMetadata(extra_fields=extra_fields)
+
+    def test_init_extra_fields_v2(self) -> None:
+        with pytest.raises(ValueError, match="Zarr format 2 group metadata does not support"):
+            GroupMetadata(zarr_format=2, extra_fields={"my_extension": {"must_understand": False}})
+
+    @pytest.mark.filterwarnings("ignore:Consolidated metadata is:zarr.errors.ZarrUserWarning")
+    def test_extra_fields_survive_rewrite(self) -> None:
+        extension = {"must_understand": False}
+        doc = {"zarr_format": 3, "node_type": "group", "attributes": {}, "ext": extension}
+        store = MemoryStore()
+        for key in ("zarr.json", "child/zarr.json"):
+            buffer = default_buffer_prototype().buffer.from_bytes(json.dumps(doc).encode())
+            sync(store.set(key, buffer))
+
+        group = zarr.open_group(store=store, mode="r+")
+        group["child"].attrs["key"] = "value"
+        zarr.consolidate_metadata(store)
+
+        written = json.loads(
+            sync(store.get("child/zarr.json", prototype=default_buffer_prototype())).to_bytes()
+        )
+        assert written == {**doc, "attributes": {"key": "value"}}
+        consolidated = zarr.open_consolidated(store)
+        assert consolidated.metadata.extra_fields == {"ext": extension}
+        assert consolidated["child"].metadata.extra_fields == {"ext": extension}
+        assert consolidated["child"].attrs["key"] == "value"
+
 
 class TestInfo:
     def test_info(self):

tests/test_metadata/test_v3.py

@@ -312,6 +312,23 @@ def test_array_metadata_extra_fields_rejected(case: ExpectFail[dict[str, Any]])
         ArrayV3Metadata.from_dict(case.input)
 
 
+@pytest.mark.parametrize("value", [{"must_understand": True}, {}, 42])
+def test_init_extra_fields_disallowed(value: object) -> None:
+    """Extra field values that are not objects with `must_understand: false` are rejected."""
+    with pytest.raises(MetadataValidationError, match="disallowed extra fields"):
+        ArrayV3Metadata(
+            shape=(10,),
+            data_type=UInt8(),
+            chunk_grid={"name": "regular", "configuration": {"chunk_shape": (10,)}},
+            chunk_key_encoding={"name": "default", "configuration": {"separator": "/"}},
+            fill_value=0,
+            codecs=({"name": "bytes", "configuration": {"endian": "little"}},),
+            attributes={},
+            dimension_names=None,
+            extra_fields={"my_ext": value},  # type: ignore[dict-item]
+        )
+
+
 def test_init_extra_fields_collision() -> None:
     """Extra field keys that collide with reserved metadata field names are rejected."""
     extra_fields: dict[str, object] = {"shape": (10,), "data_type": "uint8"}