feature(group): add group setitem api

Author: jhamman

src/zarr/api/asynchronous.py

@@ -396,12 +396,16 @@ async def save_array(
 
     mode = kwargs.pop("mode", None)
     store_path = await make_store_path(store, path=path, mode=mode, storage_options=storage_options)
+    if np.isscalar(arr):
+        arr = np.array(arr)
+    shape = arr.shape
+    chunks = getattr(arr, "chunks", shape)  # for array-likes with chunks attribute
     new = await AsyncArray.create(
         store_path,
         zarr_format=zarr_format,
-        shape=arr.shape,
+        shape=shape,
         dtype=arr.dtype,
-        chunks=arr.shape,
+        chunks=chunks,
         **kwargs,
     )
     await new.setitem(slice(None), arr)

src/zarr/core/group.py

@@ -596,6 +596,12 @@ def from_dict(
             store_path=store_path,
         )
 
+    async def setitem(self, key: str, value: Any) -> None:
+        path = self.store_path / key
+        await async_api.save_array(
+            store=path, arr=value, zarr_format=self.metadata.zarr_format, exists_ok=True
+        )
+
     async def getitem(
         self,
         key: str,
@@ -1369,8 +1375,8 @@ def __len__(self) -> int:
         return self.nmembers()
 
     def __setitem__(self, key: str, value: Any) -> None:
-        """__setitem__ is not supported in v3"""
-        raise NotImplementedError
+        """Create a new array"""
+        self._sync(self._async_group.setitem(key, value))
 
     def __repr__(self) -> str:
         return f"<Group {self.store_path}>"

src/zarr/storage/zip.py

@@ -219,7 +219,10 @@ async def set_if_not_exists(self, key: str, value: Buffer) -> None:
 
     async def delete(self, key: str) -> None:
         # docstring inherited
-        raise NotImplementedError
+        # we choose to only raise NotImplementedError here if the key exists
+        # this allows the array/group APIs to avoid the overhead of existence checks
+        if await self.exists(key):
+            raise NotImplementedError
 
     async def exists(self, key: str) -> bool:
         # docstring inherited

tests/v3/test_group.py

@@ -391,8 +391,29 @@ def test_group_setitem(store: Store, zarr_format: ZarrFormat) -> None:
     Test the `Group.__setitem__` method.
     """
     group = Group.from_store(store, zarr_format=zarr_format)
-    with pytest.raises(NotImplementedError):
-        group["key"] = 10
+    arr = np.ones((2, 4))
+    group["key"] = arr
+    assert group["key"].shape == (2, 4)
+    np.testing.assert_array_equal(group["key"][:], arr)
+
+    if store.supports_deletes:
+        key = "key"
+    else:
+        # overwriting with another array requires deletes
+        # for stores that don't support this, we just use a new key
+        key = "key2"
+
+    # overwrite with another array
+    arr = np.zeros((3, 5))
+    group[key] = arr
+    assert group[key].shape == (3, 5)
+    np.testing.assert_array_equal(group[key], arr)
+
+    # overwrite with a scalar
+    # separate bug!
+    # group["key"] = 1.5
+    # assert group["key"].shape == ()
+    # assert group["key"][:] == 1
 
 
 def test_group_contains(store: Store, zarr_format: ZarrFormat) -> None: