Skip to content

Commit

Permalink
feat: functional .children method for groups
Browse files Browse the repository at this point in the history
  • Loading branch information
d-v-b committed Mar 26, 2024
1 parent 76c3450 commit e492be2
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 4 deletions.
37 changes: 34 additions & 3 deletions src/zarr/v3/group.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,17 @@
import asyncio
import json
import logging
from typing import Any, Dict, Literal, Optional, Union, AsyncIterator, Iterator, List
from typing import (
Any,
AsyncGenerator,
Dict,
Literal,
Optional,
Union,
AsyncIterator,
Iterator,
List,
)
from zarr.v3.abc.metadata import Metadata

from zarr.v3.array import AsyncArray, Array
Expand Down Expand Up @@ -271,8 +281,29 @@ def __repr__(self):
async def nchildren(self) -> int:
raise NotImplementedError

async def children(self) -> AsyncIterator[AsyncArray, AsyncGroup]:
raise NotImplementedError
async def children(self) -> AsyncGenerator[AsyncArray, AsyncGroup]:
"""
Returns an async iterator over the arrays and groups contained in this group.
"""
if not self.store_path.store.supports_listing:
msg = (
f"The store associated with this group ({type(self.store_path.store)}) "
"does not support listing, "
"specifically the `list_dir` method. "
"This function requires a store that supports listing."
)

raise ValueError(msg)
subkeys = await self.store_path.store.list_dir(self.store_path.path)
# would be nice to make these special keys accessible programmatically,
# and scoped to specific zarr versions
subkeys_filtered = filter(lambda v: v not in ("zarr.json", ".zgroup", ".zattrs"), subkeys)
# might be smarter to wrap this in asyncio gather
for subkey in subkeys_filtered:
try:
yield await self.getitem(subkey)
except ValueError:
pass

async def contains(self, child: str) -> bool:
raise NotImplementedError
Expand Down
4 changes: 3 additions & 1 deletion tests/test_group_v3.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,13 +21,15 @@ def test_group(store_path) -> None:
runtime_configuration=RuntimeConfiguration(),
)
group = Group(agroup)

assert agroup.metadata is group.metadata

# create two groups
foo = group.create_group("foo")
bar = foo.create_group("bar", attributes={"baz": "qux"})

# check that bar is in the children of foo
assert foo.children == [bar]

# create an array from the "bar" group
data = np.arange(0, 4 * 4, dtype="uint16").reshape((4, 4))
arr = bar.create_array(
Expand Down

0 comments on commit e492be2

Please sign in to comment.