diff --git a/changes/4451.bugfix.md b/changes/4451.bugfix.md new file mode 100644 index 0000000000..7d4c00a0f9 --- /dev/null +++ b/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`. diff --git a/src/zarr/core/group.py b/src/zarr/core/group.py index d734e6b7cd..304cab10b8 100644 --- a/src/zarr/core/group.py +++ b/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: diff --git a/src/zarr/core/metadata/v3.py b/src/zarr/core/metadata/v3.py index 11f3eb593d..2959cadd74 100644 --- a/src/zarr/core/metadata/v3.py +++ b/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) diff --git a/src/zarr/core/sync_group.py b/src/zarr/core/sync_group.py index 8af514e938..6dcd5b3271 100644 --- a/src/zarr/core/sync_group.py +++ b/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) diff --git a/tests/test_group.py b/tests/test_group.py index 31fbd138cd..0dfa6ac922 100644 --- a/tests/test_group.py +++ b/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): diff --git a/tests/test_metadata/test_v3.py b/tests/test_metadata/test_v3.py index 9f78ed7b70..be2564a570 100644 --- a/tests/test_metadata/test_v3.py +++ b/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"}