fix(attrs): update array attrs in place (zarr-developers#2329)

* fix(attrs): update array attrs in place

* improve test

Author: jhamman

src/zarr/core/array.py

@@ -896,11 +896,14 @@ async def _delete_key(key: str) -> None:
         return replace(self, metadata=new_metadata)
 
     async def update_attributes(self, new_attributes: dict[str, JSON]) -> Self:
-        new_metadata = self.metadata.update_attributes(new_attributes)
+        # 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
-        await self._save_metadata(new_metadata)
-        return replace(self, metadata=new_metadata)
+        await self._save_metadata(self.metadata)
+
+        return self
 
     def __repr__(self) -> str:
         return f"<AsyncArray {self.store_path} shape={self.shape} dtype={self.dtype}>"

tests/v3/test_array.py

@@ -406,3 +406,15 @@ def test_vlen_errors() -> None:
             dtype="<U4",
             codecs=[BytesCodec(), VLenBytesCodec()],
         )
+
+
+@pytest.mark.parametrize("zarr_format", [2, 3])
+def test_update_attrs(zarr_format: int) -> None:
+    # regression test for https://github.com/zarr-developers/zarr-python/issues/2328
+    store = MemoryStore({}, mode="w")
+    arr = Array.create(store=store, shape=5, chunk_shape=5, dtype="f8", zarr_format=zarr_format)
+    arr.attrs["foo"] = "bar"
+    assert arr.attrs["foo"] == "bar"
+
+    arr2 = zarr.open_array(store=store, zarr_format=zarr_format)
+    assert arr2.attrs["foo"] == "bar"

tests/v3/test_group.py

@@ -1257,3 +1257,12 @@ def test_from_dict_extra_fields(self):
         result = GroupMetadata.from_dict(data)
         expected = GroupMetadata(attributes={"key": "value"}, zarr_format=2)
         assert result == expected
+
+
+def test_update_attrs() -> None:
+    # regression test for https://github.com/zarr-developers/zarr-python/issues/2328
+    root = Group.from_store(
+        MemoryStore({}, mode="w"),
+    )
+    root.attrs["foo"] = "bar"
+    assert root.attrs["foo"] == "bar"