diff --git a/changes/199.feature.md b/changes/199.feature.md new file mode 100644 index 0000000000..1af0fa300a --- /dev/null +++ b/changes/199.feature.md @@ -0,0 +1 @@ +Added `replace_attributes` method to `Array`, `AsyncArray`, `Group`, and `AsyncGroup` to overwrite all attributes (as opposed to `update_attributes`, which merges). diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 2c2a4622e3..8324e77039 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -1703,10 +1703,42 @@ async def update_attributes(self, new_attributes: dict[str, JSON]) -> Self: Notes ----- - This method is asynchronous and should be awaited. - - The updated attributes will be merged with existing attributes, and any conflicts will be - overwritten by the new values. + - `new_attributes` is MERGED into the existing attributes: keys that are + already present and not in `new_attributes` are preserved, and conflicting + keys are overwritten by the new values. To drop keys, use + `replace_attributes` instead. """ - await _update_attributes(self, new_attributes) + await _update_attributes(self, new_attributes, merge=True) + return self + + async def replace_attributes(self, new_attributes: dict[str, JSON]) -> Self: + """ + Asynchronously replace all of the array's attributes. + + Parameters + ---------- + new_attributes : dict of str to JSON + A dictionary of attributes that will replace the array's existing attributes. + The keys represent attribute names, and the values must be JSON-compatible. + + Returns + ------- + AsyncArray + The array with the replaced attributes. + + Raises + ------ + ValueError + If the attributes are invalid or incompatible with the array's metadata. + + Notes + ----- + - This method is asynchronous and should be awaited. + - This REPLACES all existing attributes with `new_attributes`. Any keys not + present in `new_attributes` are removed. To merge into the existing + attributes instead, use `update_attributes`. + """ + await _update_attributes(self, new_attributes, merge=False) return self def __repr__(self) -> str: @@ -3884,12 +3916,43 @@ def update_attributes(self, new_attributes: dict[str, JSON]) -> Self: Notes ----- - - The updated attributes will be merged with existing attributes, and any conflicts will be - overwritten by the new values. + - `new_attributes` is MERGED into the existing attributes: keys that are + already present and not in `new_attributes` are preserved, and conflicting + keys are overwritten by the new values. To drop keys, use + `replace_attributes` instead. """ new_array = sync(self.async_array.update_attributes(new_attributes)) return type(self)(new_array) + def replace_attributes(self, new_attributes: dict[str, JSON]) -> Self: + """ + Replace all of the array's attributes. + + Parameters + ---------- + new_attributes : dict + A dictionary of attributes that will replace the array's existing attributes. + The keys represent attribute names, and the values must be JSON-compatible. + + Returns + ------- + Array + The array with the replaced attributes. + + Raises + ------ + ValueError + If the attributes are invalid or incompatible with the array's metadata. + + Notes + ----- + - This REPLACES all existing attributes with `new_attributes`. Any keys not + present in `new_attributes` are removed. To merge into the existing + attributes instead, use `update_attributes`. + """ + new_array = sync(self.async_array.replace_attributes(new_attributes)) + return type(self)(new_array) + def __repr__(self) -> str: return f"" @@ -5990,26 +6053,41 @@ async def _append( async def _update_attributes( array: AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata], new_attributes: dict[str, JSON], + *, + merge: bool = True, ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: """ - Update the array's attributes. + Update or replace the array's attributes. Parameters ---------- array : AsyncArray The array whose attributes to update. new_attributes : dict[str, JSON] - A dictionary of new attributes to update or add to the array. + A dictionary of new attributes for the array. + merge : bool, default True + If True, merge `new_attributes` into the existing attributes. If False, + replace all existing attributes with `new_attributes`. Returns ------- AsyncArray The array with the updated attributes. """ - array.metadata.attributes.update(new_attributes) + if merge: + attributes = {**array.metadata.attributes, **new_attributes} + else: + attributes = dict(new_attributes) - # Write new metadata - await save_metadata(array.store_path, array.metadata) + # Build new metadata without mutating the existing (frozen) metadata. + new_metadata = array.metadata.update_attributes(attributes) + + # Write new metadata to the store first, so a failed write does not corrupt + # the in-memory state. + await save_metadata(array.store_path, new_metadata) + + # Only swap the in-memory metadata once the write has succeeded. + object.__setattr__(array, "metadata", new_metadata) return array diff --git a/src/zarr/core/attributes.py b/src/zarr/core/attributes.py index 7f29e44365..b10256835a 100644 --- a/src/zarr/core/attributes.py +++ b/src/zarr/core/attributes.py @@ -50,8 +50,7 @@ def put(self, d: dict[str, JSON]) -> None: #> {'a': '3', 'c': 4} ``` """ - self._obj.metadata.attributes.clear() - self._obj = self._obj.update_attributes(d) + self._obj = self._obj.replace_attributes(d) def asdict(self) -> dict[str, JSON]: return dict(self._obj.metadata.attributes) diff --git a/src/zarr/core/group.py b/src/zarr/core/group.py index 52eaa3e144..aec4411627 100644 --- a/src/zarr/core/group.py +++ b/src/zarr/core/group.py @@ -1204,24 +1204,65 @@ async def require_array( return ds + async def _set_attributes(self, attributes: dict[str, Any]) -> AsyncGroup: + """Write `attributes` as the group's attributes, atomically. + + Builds new metadata without mutating the existing (frozen) metadata, writes + it to the store first, and only swaps the in-memory metadata once the write + has succeeded. + """ + new_metadata = replace(self.metadata, attributes=attributes) + + # Write new metadata to the store first, so a failed write does not corrupt + # the in-memory state. + await save_metadata(self.store_path, new_metadata) + + # Only swap the in-memory metadata once the write has succeeded. + object.__setattr__(self, "metadata", new_metadata) + + return self + async def update_attributes(self, new_attributes: dict[str, Any]) -> AsyncGroup: """Update group attributes. Parameters ---------- new_attributes : dict - New attributes to set on the group. + New attributes to merge into the group's existing attributes. Returns ------- self : AsyncGroup + + Notes + ----- + - `new_attributes` is MERGED into the existing attributes: keys that are + already present and not in `new_attributes` are preserved, and conflicting + keys are overwritten by the new values. To drop keys, use + `replace_attributes` instead. """ - self.metadata.attributes.update(new_attributes) + merged = {**self.metadata.attributes, **new_attributes} + return await self._set_attributes(merged) - # Write new metadata - await self._save_metadata() + async def replace_attributes(self, new_attributes: dict[str, Any]) -> AsyncGroup: + """Replace all of the group's attributes. - return self + Parameters + ---------- + new_attributes : dict + Attributes that will replace the group's existing attributes. + + Returns + ------- + self : AsyncGroup + + Notes + ----- + - This REPLACES all existing attributes with `new_attributes`. Any keys not + present in `new_attributes` are removed. To merge into the existing + attributes instead, use `update_attributes`. + """ + return await self._set_attributes(dict(new_attributes)) def __repr__(self) -> str: return f"" @@ -2055,6 +2096,11 @@ def synchronizer(self) -> None: def update_attributes(self, new_attributes: dict[str, Any]) -> Group: """Update the attributes of this group. + `new_attributes` is MERGED into the existing attributes: keys that are + already present and not in `new_attributes` are preserved, and conflicting + keys are overwritten by the new values. To drop keys, use + `replace_attributes` instead. + Examples -------- >>> import zarr @@ -2065,6 +2111,24 @@ def update_attributes(self, new_attributes: dict[str, Any]) -> Group: self._sync(self._async_group.update_attributes(new_attributes)) return self + def replace_attributes(self, new_attributes: dict[str, Any]) -> Group: + """Replace all of the attributes of this group. + + This REPLACES all existing attributes with `new_attributes`. Any keys not + present in `new_attributes` are removed. To merge into the existing + attributes instead, use `update_attributes`. + + Examples + -------- + >>> import zarr + >>> group = zarr.group(attributes={"foo": "bar"}) + >>> group = group.replace_attributes({"baz": "qux"}) + >>> group.attrs.asdict() + {'baz': 'qux'} + """ + self._sync(self._async_group.replace_attributes(new_attributes)) + return self + def nmembers(self, max_depth: int | None = 0) -> int: """Count the number of members in this group. diff --git a/tests/test_attributes.py b/tests/test_attributes.py index 269704d2a0..3664cf828a 100644 --- a/tests/test_attributes.py +++ b/tests/test_attributes.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import json from typing import TYPE_CHECKING, Any @@ -8,9 +10,9 @@ import zarr.core.attributes import zarr.storage from tests.conftest import deep_nan_equal -from zarr.core.common import ZarrFormat if TYPE_CHECKING: + from zarr.core.common import ZarrFormat from zarr.types import AnyArray @@ -74,6 +76,94 @@ def test_update_no_changes() -> None: assert dict(z.attrs) == {"a": [], "b": 3} +def _make_node( + store: zarr.storage.MemoryStore, + *, + group: bool, + attributes: dict[str, Any] | None = None, +) -> zarr.Group | AnyArray: + """Create a fresh sync group or array with the given attributes.""" + if group: + return zarr.create_group(store, attributes=attributes) + return zarr.create_array( + store=store, shape=10, dtype=int, attributes=attributes, overwrite=True + ) + + +@pytest.mark.parametrize("group", [True, False]) +def test_update_attributes_merges(group: bool) -> None: + """`update_attributes` merges into existing attributes; pre-existing keys survive.""" + store = zarr.storage.MemoryStore() + z = _make_node(store, group=group, attributes={"a": 1}) + z.update_attributes({"b": 2}) + assert dict(z.attrs) == {"a": 1, "b": 2} + + +@pytest.mark.parametrize("group", [True, False]) +def test_update_attributes_no_mutation(group: bool) -> None: + """`update_attributes` must not mutate the original frozen metadata object.""" + store = zarr.storage.MemoryStore() + z = _make_node(store, group=group, attributes={"a": 1}) + old_metadata = z.metadata + snapshot = dict(old_metadata.attributes) + z.update_attributes({"b": 2}) + assert dict(old_metadata.attributes) == snapshot + + +@pytest.mark.parametrize("group", [True, False]) +def test_replace_attributes_replaces_sync(group: bool) -> None: + """`replace_attributes` drops keys absent from the new dict (sync).""" + store = zarr.storage.MemoryStore() + z = _make_node(store, group=group, attributes={"a": 1, "b": 2}) + z.replace_attributes({"a": 3, "c": 4}) + assert dict(z.attrs) == {"a": 3, "c": 4} + + +@pytest.mark.parametrize("group", [True, False]) +async def test_replace_attributes_replaces_async(group: bool) -> None: + """`replace_attributes` drops keys absent from the new dict (async).""" + store = zarr.storage.MemoryStore() + z = _make_node(store, group=group, attributes={"a": 1, "b": 2}) + async_obj = z._async_group if isinstance(z, zarr.Group) else z.async_array + await async_obj.replace_attributes({"a": 3, "c": 4}) + assert dict(async_obj.metadata.attributes) == {"a": 3, "c": 4} + + +@pytest.mark.parametrize("group", [True, False]) +def test_replace_attributes_no_mutation(group: bool) -> None: + """`replace_attributes` must not mutate the original frozen metadata object.""" + store = zarr.storage.MemoryStore() + z = _make_node(store, group=group, attributes={"a": 1, "b": 2}) + old_metadata = z.metadata + snapshot = dict(old_metadata.attributes) + z.replace_attributes({"a": 3, "c": 4}) + assert dict(old_metadata.attributes) == snapshot + + +@pytest.mark.parametrize("group", [True, False]) +def test_put_replaces(group: bool) -> None: + """`attrs.put` replaces all attributes, dropping absent keys.""" + store = zarr.storage.MemoryStore() + z = _make_node(store, group=group, attributes={"a": 1, "b": 2}) + z.attrs.put({"a": 3, "c": 4}) + assert dict(z.attrs) == {"a": 3, "c": 4} + + +@pytest.mark.parametrize("group", [True, False]) +def test_replace_attributes_persists(group: bool) -> None: + """After replace, reopening from the store reflects the dropped keys.""" + store = zarr.storage.MemoryStore() + z = _make_node(store, group=group, attributes={"a": 1, "b": 2}) + z.replace_attributes({"a": 3, "c": 4}) + + z2: zarr.Group | AnyArray + if group: + z2 = zarr.open_group(store) + else: + z2 = zarr.open_array(store) + assert dict(z2.attrs) == {"a": 3, "c": 4} + + @pytest.mark.parametrize("group", [True, False]) def test_del_works(group: bool) -> None: store = zarr.storage.MemoryStore()