Prevent update_attributes from erasing all prior attributes (zarr-developers#2870)

* Prevent update_attributes from erasing all prior attributes

* Fix incorrect group attribute update tests

* Add changelog entry

Author: moradology

changes/2870.bugfix.rst

@@ -0,0 +1 @@
+Prevent update_attributes calls from deleting old attributes

src/zarr/core/array.py

@@ -1609,8 +1609,6 @@ async def update_attributes(self, new_attributes: dict[str, JSON]) -> Self:
         - The updated attributes will be merged with existing attributes, and any conflicts will be
           overwritten by the new values.
         """
-        # metadata.attributes is "frozen" so we simply clear and update the dict
-        self.metadata.attributes.clear()
         self.metadata.attributes.update(new_attributes)
 
         # Write new metadata

src/zarr/core/attributes.py

@@ -50,6 +50,7 @@ def put(self, d: dict[str, JSON]) -> None:
            >>> attrs
            {'a': 3, 'c': 4}
         """
+        self._obj.metadata.attributes.clear()
         self._obj = self._obj.update_attributes(d)
 
     def asdict(self) -> dict[str, JSON]:

src/zarr/core/common.py

@@ -4,7 +4,7 @@
 import functools
 import operator
 import warnings
-from collections.abc import Iterable, Mapping
+from collections.abc import Iterable, Mapping, Sequence
 from enum import Enum
 from itertools import starmap
 from typing import (
@@ -37,7 +37,7 @@
 ChunkCoordsLike = Iterable[int]
 ZarrFormat = Literal[2, 3]
 NodeType = Literal["array", "group"]
-JSON = str | int | float | Mapping[str, "JSON"] | tuple["JSON", ...] | None
+JSON = str | int | float | Mapping[str, "JSON"] | Sequence["JSON"] | None
 MemoryOrder = Literal["C", "F"]
 AccessModeLiteral = Literal["r", "r+", "a", "w", "w-"]
 

src/zarr/core/group.py

@@ -1254,8 +1254,6 @@ async def update_attributes(self, new_attributes: dict[str, Any]) -> AsyncGroup:
         -------
         self : AsyncGroup
         """
-        # metadata.attributes is "frozen" so we simply clear and update the dict
-        self.metadata.attributes.clear()
         self.metadata.attributes.update(new_attributes)
 
         # Write new metadata

tests/test_attributes.py

@@ -20,3 +20,42 @@ def test_asdict() -> None:
     )
     result = attrs.asdict()
     assert result == {"a": 1, "b": 2}
+
+
+def test_update_attributes_preserves_existing() -> None:
+    """
+    Test that `update_attributes` only updates the specified attributes
+    and preserves existing ones.
+    """
+    store = zarr.storage.MemoryStore()
+    z = zarr.create(10, store=store, overwrite=True)
+    z.attrs["a"] = []
+    z.attrs["b"] = 3
+    assert dict(z.attrs) == {"a": [], "b": 3}
+
+    z.update_attributes({"a": [3, 4], "c": 4})
+    assert dict(z.attrs) == {"a": [3, 4], "b": 3, "c": 4}
+
+
+def test_update_empty_attributes() -> None:
+    """
+    Ensure updating when initial attributes are empty works.
+    """
+    store = zarr.storage.MemoryStore()
+    z = zarr.create(10, store=store, overwrite=True)
+    assert dict(z.attrs) == {}
+    z.update_attributes({"a": [3, 4], "c": 4})
+    assert dict(z.attrs) == {"a": [3, 4], "c": 4}
+
+
+def test_update_no_changes() -> None:
+    """
+    Ensure updating when no new or modified attributes does not alter existing ones.
+    """
+    store = zarr.storage.MemoryStore()
+    z = zarr.create(10, store=store, overwrite=True)
+    z.attrs["a"] = []
+    z.attrs["b"] = 3
+
+    z.update_attributes({})
+    assert dict(z.attrs) == {"a": [], "b": 3}

tests/test_group.py

@@ -604,7 +604,10 @@ def test_group_update_attributes(store: Store, zarr_format: ZarrFormat) -> None:
     assert group.attrs == attrs
     new_attrs = {"bar": 100}
     new_group = group.update_attributes(new_attrs)
-    assert new_group.attrs == new_attrs
+
+    updated_attrs = attrs.copy()
+    updated_attrs.update(new_attrs)
+    assert new_group.attrs == updated_attrs
 
 
 async def test_group_update_attributes_async(store: Store, zarr_format: ZarrFormat) -> None:
@@ -1008,7 +1011,9 @@ async def test_asyncgroup_update_attributes(store: Store, zarr_format: ZarrForma
     )
 
     agroup_new_attributes = await agroup.update_attributes(attributes_new)
-    assert agroup_new_attributes.attrs == attributes_new
+    attributes_updated = attributes_old.copy()
+    attributes_updated.update(attributes_new)
+    assert agroup_new_attributes.attrs == attributes_updated
 
 
 @pytest.mark.parametrize("store", ["local"], indirect=["store"])