Skip to content

Commit f4af51c

Browse files
authored
feature(group): add group setitem api (#2393)
* feature(group): add group setitem api * arrays proxy * rollback to simple version * rollback deprecation * rollback ... * Update tests/test_group.py
1 parent 109f71f commit f4af51c

File tree

4 files changed

+51
-7
lines changed

4 files changed

+51
-7
lines changed

src/zarr/api/asynchronous.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -396,12 +396,16 @@ async def save_array(
396396

397397
mode = kwargs.pop("mode", None)
398398
store_path = await make_store_path(store, path=path, mode=mode, storage_options=storage_options)
399+
if np.isscalar(arr):
400+
arr = np.array(arr)
401+
shape = arr.shape
402+
chunks = getattr(arr, "chunks", None) # for array-likes with chunks attribute
399403
new = await AsyncArray.create(
400404
store_path,
401405
zarr_format=zarr_format,
402-
shape=arr.shape,
406+
shape=shape,
403407
dtype=arr.dtype,
404-
chunks=arr.shape,
408+
chunks=chunks,
405409
**kwargs,
406410
)
407411
await new.setitem(slice(None), arr)

src/zarr/core/group.py

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -600,6 +600,23 @@ def from_dict(
600600
store_path=store_path,
601601
)
602602

603+
async def setitem(self, key: str, value: Any) -> None:
604+
"""Fastpath for creating a new array
605+
606+
New arrays will be created with default array settings for the array type.
607+
608+
Parameters
609+
----------
610+
key : str
611+
Array name
612+
value : array-like
613+
Array data
614+
"""
615+
path = self.store_path / key
616+
await async_api.save_array(
617+
store=path, arr=value, zarr_format=self.metadata.zarr_format, exists_ok=True
618+
)
619+
603620
async def getitem(
604621
self,
605622
key: str,
@@ -1394,8 +1411,11 @@ def __len__(self) -> int:
13941411
return self.nmembers()
13951412

13961413
def __setitem__(self, key: str, value: Any) -> None:
1397-
"""__setitem__ is not supported in v3"""
1398-
raise NotImplementedError
1414+
"""Fastpath for creating a new array.
1415+
1416+
New arrays will be created using default settings for the array type.
1417+
"""
1418+
self._sync(self._async_group.setitem(key, value))
13991419

14001420
def __repr__(self) -> str:
14011421
return f"<Group {self.store_path}>"

src/zarr/storage/zip.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -219,7 +219,10 @@ async def set_if_not_exists(self, key: str, value: Buffer) -> None:
219219

220220
async def delete(self, key: str) -> None:
221221
# docstring inherited
222-
raise NotImplementedError
222+
# we choose to only raise NotImplementedError here if the key exists
223+
# this allows the array/group APIs to avoid the overhead of existence checks
224+
if await self.exists(key):
225+
raise NotImplementedError
223226

224227
async def exists(self, key: str) -> bool:
225228
# docstring inherited

tests/test_group.py

Lines changed: 19 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -427,8 +427,25 @@ def test_group_setitem(store: Store, zarr_format: ZarrFormat) -> None:
427427
Test the `Group.__setitem__` method.
428428
"""
429429
group = Group.from_store(store, zarr_format=zarr_format)
430-
with pytest.raises(NotImplementedError):
431-
group["key"] = 10
430+
arr = np.ones((2, 4))
431+
group["key"] = arr
432+
assert list(group.array_keys()) == ["key"]
433+
assert group["key"].shape == (2, 4)
434+
np.testing.assert_array_equal(group["key"][:], arr)
435+
436+
if store.supports_deletes:
437+
key = "key"
438+
else:
439+
# overwriting with another array requires deletes
440+
# for stores that don't support this, we just use a new key
441+
key = "key2"
442+
443+
# overwrite with another array
444+
arr = np.zeros((3, 5))
445+
group[key] = arr
446+
assert key in list(group.array_keys())
447+
assert group[key].shape == (3, 5)
448+
np.testing.assert_array_equal(group[key], arr)
432449

433450

434451
def test_group_contains(store: Store, zarr_format: ZarrFormat) -> None:

0 commit comments

Comments
 (0)