Skip to content

Commit cb93ef8

Browse files
feat: add type-safe get_array and get_group methods to AsyncGroup and Group (#4128)
* chore(deps): bump the actions group across 1 directory with 8 updates (#176) Bumps the actions group with 8 updates in the / directory: | Package | From | To | | --- | --- | --- | | [prefix-dev/setup-pixi](https://github.com/prefix-dev/setup-pixi) | `0.9.5` | `0.9.6` | | [codecov/codecov-action](https://github.com/codecov/codecov-action) | `6.0.0` | `6.0.1` | | [github/issue-metrics](https://github.com/github/issue-metrics) | `4.2.2` | `4.2.7` | | [j178/prek-action](https://github.com/j178/prek-action) | `2.0.3` | `2.0.4` | | [actions/upload-artifact](https://github.com/actions/upload-artifact) | `7.0.0` | `7.0.1` | | [actions/download-artifact](https://github.com/actions/download-artifact) | `7.0.0` | `8.0.1` | | [pypa/gh-action-pypi-publish](https://github.com/pypa/gh-action-pypi-publish) | `1.13.0` | `1.14.0` | | [zizmorcore/zizmor-action](https://github.com/zizmorcore/zizmor-action) | `0.5.3` | `0.5.6` | Updates `prefix-dev/setup-pixi` from 0.9.5 to 0.9.6 - [Release notes](https://github.com/prefix-dev/setup-pixi/releases) - [Commits](prefix-dev/setup-pixi@1b2de7f...5185adf) Updates `codecov/codecov-action` from 6.0.0 to 6.0.1 - [Release notes](https://github.com/codecov/codecov-action/releases) - [Changelog](https://github.com/codecov/codecov-action/blob/main/CHANGELOG.md) - [Commits](codecov/codecov-action@57e3a13...e79a696) Updates `github/issue-metrics` from 4.2.2 to 4.2.7 - [Release notes](https://github.com/github/issue-metrics/releases) - [Commits](github-community-projects/issue-metrics@c9e9838...1e38d5e) Updates `j178/prek-action` from 2.0.3 to 2.0.4 - [Release notes](https://github.com/j178/prek-action/releases) - [Commits](j178/prek-action@6ad8027...bdca6f1) Updates `actions/upload-artifact` from 7.0.0 to 7.0.1 - [Release notes](https://github.com/actions/upload-artifact/releases) - [Commits](actions/upload-artifact@v7...043fb46) Updates `actions/download-artifact` from 7.0.0 to 8.0.1 - [Release notes](https://github.com/actions/download-artifact/releases) - [Commits](actions/download-artifact@v7...3e5f45b) Updates `pypa/gh-action-pypi-publish` from 1.13.0 to 1.14.0 - [Release notes](https://github.com/pypa/gh-action-pypi-publish/releases) - [Commits](pypa/gh-action-pypi-publish@v1.13.0...cef2210) Updates `zizmorcore/zizmor-action` from 0.5.3 to 0.5.6 - [Release notes](https://github.com/zizmorcore/zizmor-action/releases) - [Commits](zizmorcore/zizmor-action@b1d7e1f...5f14fd0) --- updated-dependencies: - dependency-name: prefix-dev/setup-pixi dependency-version: 0.9.6 dependency-type: direct:production update-type: version-update:semver-patch dependency-group: actions - dependency-name: codecov/codecov-action dependency-version: 6.0.1 dependency-type: direct:production update-type: version-update:semver-patch dependency-group: actions - dependency-name: github/issue-metrics dependency-version: 4.2.7 dependency-type: direct:production update-type: version-update:semver-patch dependency-group: actions - dependency-name: j178/prek-action dependency-version: 2.0.4 dependency-type: direct:production update-type: version-update:semver-patch dependency-group: actions - dependency-name: actions/upload-artifact dependency-version: 7.0.1 dependency-type: direct:production update-type: version-update:semver-patch dependency-group: actions - dependency-name: actions/download-artifact dependency-version: 8.0.1 dependency-type: direct:production update-type: version-update:semver-major dependency-group: actions - dependency-name: pypa/gh-action-pypi-publish dependency-version: 1.14.0 dependency-type: direct:production update-type: version-update:semver-minor dependency-group: actions - dependency-name: zizmorcore/zizmor-action dependency-version: 0.5.6 dependency-type: direct:production update-type: version-update:semver-patch dependency-group: actions ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> * feat: add type-safe get_array and get_group methods to AsyncGroup and Group Add `get_array` and `get_group` methods that return the child node at a given path or raise: `ArrayNotFoundError`/`GroupNotFoundError` when no node exists, and `ContainsGroupError`/`ContainsArrayError` when the node is the wrong kind. Paths mirror `getitem` semantics, so nested paths like "subgroup/subarray" work on both the plain and consolidated-metadata lookup routes. Existing tests that fetched a child via `getitem` and then manually narrowed the type with `isinstance` asserts or a walrus expression now use the new methods instead. Assisted-by: ClaudeCode:claude-fable-5 * fix: pass pre-formatted messages to error constructors in get_array/get_group The multi-argument template form of BaseZarrError.__init__ is documented as deprecated; build the message string at the raise site instead, matching every other call site in the codebase. Assisted-by: ClaudeCode:claude-fable-5 * docs: add changelog entry for get_array/get_group Assisted-by: ClaudeCode:claude-fable-5 * docs: demonstrate get_array/get_group in the groups user guide Assisted-by: ClaudeCode:claude-fable-5 --------- Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
1 parent 0b72757 commit cb93ef8

6 files changed

Lines changed: 239 additions & 13 deletions

File tree

changes/4128.feature.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Added `Group.get_array`, `Group.get_group`, `AsyncGroup.get_array`, and `AsyncGroup.get_group`: type-safe accessors that return the child array or group at a given path, raising `ArrayNotFoundError` / `GroupNotFoundError` if no node exists there, and `ContainsGroupError` / `ContainsArrayError` if the node is not of the requested kind. Unlike `Group.__getitem__`, which returns `Array | Group`, these methods have precise return types. Nested paths like `"subgroup/subarray"` are supported.

docs/user-guide/groups.md

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,29 @@ print(root['foo/bar'])
5151
print(root['foo/bar/baz'])
5252
```
5353

54+
Accessing a member with `[]` returns either an [`zarr.Array`][] or a [`zarr.Group`][], depending on
55+
what is stored at the given path. When you expect a node of a particular kind, use
56+
[`zarr.Group.get_array`][] or [`zarr.Group.get_group`][] instead. These methods accept the same
57+
paths as `[]`, but they have precise return types and raise an error if no node exists at the
58+
given path, or if the node is not of the expected kind:
59+
60+
```python exec="true" session="groups" source="above" result="ansi"
61+
print(root.get_group('foo'))
62+
```
63+
64+
```python exec="true" session="groups" source="above" result="ansi"
65+
print(root.get_array('foo/bar/baz'))
66+
```
67+
68+
```python exec="true" session="groups" source="above" result="ansi"
69+
from zarr.errors import ContainsGroupError
70+
71+
try:
72+
root.get_array('foo')
73+
except ContainsGroupError as e:
74+
print(e)
75+
```
76+
5477
The [`zarr.Group.tree`][] method can be used to print a tree
5578
representation of the hierarchy, e.g.:
5679

src/zarr/core/group.py

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@
5151
from zarr.core.metadata.io import save_metadata
5252
from zarr.core.sync import SyncMixin, sync
5353
from zarr.errors import (
54+
ArrayNotFoundError,
5455
ContainsArrayError,
5556
ContainsGroupError,
5657
GroupNotFoundError,
@@ -820,6 +821,70 @@ async def get[DefaultT](
820821
except KeyError:
821822
return default
822823

824+
async def get_array(self, path: str) -> AnyAsyncArray:
825+
"""Obtain an array member of this group, raising if it is absent or not an array.
826+
827+
Parameters
828+
----------
829+
path : str
830+
Path of the array relative to this group. May contain `/` to reference
831+
a member of a subgroup, e.g. `subgroup/subarray`.
832+
833+
Returns
834+
-------
835+
AsyncArray
836+
The array at the given path.
837+
838+
Raises
839+
------
840+
ArrayNotFoundError
841+
If no node exists at the given path.
842+
ContainsGroupError
843+
If the node at the given path is a group rather than an array.
844+
"""
845+
store_path = self.store_path / path
846+
try:
847+
node = await self.getitem(path)
848+
except KeyError as e:
849+
msg = f"No array found in store {store_path.store!r} at path {store_path.path!r}"
850+
raise ArrayNotFoundError(msg) from e
851+
if isinstance(node, AsyncGroup):
852+
msg = f"A group exists in store {store_path.store!r} at path {store_path.path!r}."
853+
raise ContainsGroupError(msg)
854+
return node
855+
856+
async def get_group(self, path: str) -> AsyncGroup:
857+
"""Obtain a group member of this group, raising if it is absent or not a group.
858+
859+
Parameters
860+
----------
861+
path : str
862+
Path of the group relative to this group. May contain `/` to reference
863+
a member of a subgroup, e.g. `subgroup/subsubgroup`.
864+
865+
Returns
866+
-------
867+
AsyncGroup
868+
The group at the given path.
869+
870+
Raises
871+
------
872+
GroupNotFoundError
873+
If no node exists at the given path.
874+
ContainsArrayError
875+
If the node at the given path is an array rather than a group.
876+
"""
877+
store_path = self.store_path / path
878+
try:
879+
node = await self.getitem(path)
880+
except KeyError as e:
881+
msg = f"No group found in store {store_path.store!r} at path {store_path.path!r}"
882+
raise GroupNotFoundError(msg) from e
883+
if isinstance(node, AsyncArray):
884+
msg = f"An array exists in store {store_path.store!r} at path {store_path.path!r}."
885+
raise ContainsArrayError(msg)
886+
return node
887+
823888
async def _save_metadata(self, ensure_parents: bool = False) -> None:
824889
await save_metadata(self.store_path, self.metadata, ensure_parents=ensure_parents)
825890

@@ -1880,6 +1945,74 @@ def get[DefaultT](
18801945
except KeyError:
18811946
return default
18821947

1948+
def get_array(self, path: str) -> AnyArray:
1949+
"""Obtain an array member of this group, raising if it is absent or not an array.
1950+
1951+
Parameters
1952+
----------
1953+
path : str
1954+
Path of the array relative to this group. May contain `/` to reference
1955+
a member of a subgroup, e.g. `subgroup/subarray`.
1956+
1957+
Returns
1958+
-------
1959+
Array
1960+
The array at the given path.
1961+
1962+
Raises
1963+
------
1964+
ArrayNotFoundError
1965+
If no node exists at the given path.
1966+
ContainsGroupError
1967+
If the node at the given path is a group rather than an array.
1968+
1969+
Examples
1970+
--------
1971+
```python
1972+
import zarr
1973+
from zarr.core.group import Group
1974+
group = Group.from_store(zarr.storage.MemoryStore())
1975+
group.create_array(name="subarray", shape=(10,), chunks=(10,), dtype="float64")
1976+
group.get_array("subarray")
1977+
# <Array memory://... shape=(10,) dtype=float64>
1978+
```
1979+
"""
1980+
return Array(self._sync(self._async_group.get_array(path)))
1981+
1982+
def get_group(self, path: str) -> Group:
1983+
"""Obtain a group member of this group, raising if it is absent or not a group.
1984+
1985+
Parameters
1986+
----------
1987+
path : str
1988+
Path of the group relative to this group. May contain `/` to reference
1989+
a member of a subgroup, e.g. `subgroup/subsubgroup`.
1990+
1991+
Returns
1992+
-------
1993+
Group
1994+
The group at the given path.
1995+
1996+
Raises
1997+
------
1998+
GroupNotFoundError
1999+
If no node exists at the given path.
2000+
ContainsArrayError
2001+
If the node at the given path is an array rather than a group.
2002+
2003+
Examples
2004+
--------
2005+
```python
2006+
import zarr
2007+
from zarr.core.group import Group
2008+
group = Group.from_store(zarr.storage.MemoryStore())
2009+
group.create_group(name="subgroup")
2010+
group.get_group("subgroup")
2011+
# <Group memory://...>
2012+
```
2013+
"""
2014+
return Group(self._sync(self._async_group.get_group(path)))
2015+
18832016
def __delitem__(self, key: str) -> None:
18842017
"""Delete a group member.
18852018

tests/test_group.py

Lines changed: 77 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -41,8 +41,10 @@
4141
from zarr.core.metadata.v3 import ArrayV3Metadata
4242
from zarr.core.sync import _collect_aiterator, sync
4343
from zarr.errors import (
44+
ArrayNotFoundError,
4445
ContainsArrayError,
4546
ContainsGroupError,
47+
GroupNotFoundError,
4648
MetadataValidationError,
4749
ZarrUserWarning,
4850
)
@@ -101,7 +103,7 @@ async def test_create_creates_parents(store: Store, zarr_format: ZarrFormat) ->
101103
root = await zarr.api.asynchronous.open_group(
102104
store=store,
103105
)
104-
agroup = await root.getitem("a")
106+
agroup = await root.get_group("a")
105107
assert agroup.attrs == {"key": "value"}
106108

107109
# create a child node with a couple intermediates
@@ -446,6 +448,77 @@ def test_group_get_with_default(store: Store, zarr_format: ZarrFormat) -> None:
446448
assert result.attrs["foo"] == "bar"
447449

448450

451+
def test_group_get_array(store: Store, zarr_format: ZarrFormat) -> None:
452+
"""
453+
`Group.get_array` returns the array at the given path, for both direct child names
454+
and nested paths, and the result is statically typed as an Array.
455+
"""
456+
group = Group.from_store(store, zarr_format=zarr_format)
457+
subgroup = group.create_group(name="subgroup")
458+
subarray = group.create_array(name="subarray", shape=(10,), chunks=(10,), dtype="uint8")
459+
subsubarray = subgroup.create_array(name="subarray", shape=(10,), chunks=(10,), dtype="uint8")
460+
461+
observed = group.get_array("subarray")
462+
assert isinstance(observed, Array)
463+
assert observed == subarray
464+
assert group.get_array("subgroup/subarray") == subsubarray
465+
466+
467+
def test_group_get_array_missing(store: Store, zarr_format: ZarrFormat) -> None:
468+
"""
469+
`Group.get_array` raises `ArrayNotFoundError` when no node exists at the given path.
470+
"""
471+
group = Group.from_store(store, zarr_format=zarr_format)
472+
with pytest.raises(ArrayNotFoundError, match="No array found in store"):
473+
group.get_array("missing")
474+
475+
476+
def test_group_get_array_wrong_node_type(store: Store, zarr_format: ZarrFormat) -> None:
477+
"""
478+
`Group.get_array` raises `ContainsGroupError` when the node at the given path is a
479+
group rather than an array.
480+
"""
481+
group = Group.from_store(store, zarr_format=zarr_format)
482+
group.create_group(name="subgroup")
483+
with pytest.raises(ContainsGroupError, match="A group exists in store"):
484+
group.get_array("subgroup")
485+
486+
487+
def test_group_get_group(store: Store, zarr_format: ZarrFormat) -> None:
488+
"""
489+
`Group.get_group` returns the group at the given path, for both direct child names
490+
and nested paths, and the result is statically typed as a Group.
491+
"""
492+
group = Group.from_store(store, zarr_format=zarr_format)
493+
subgroup = group.create_group(name="subgroup")
494+
subsubgroup = subgroup.create_group(name="subsubgroup")
495+
496+
observed = group.get_group("subgroup")
497+
assert isinstance(observed, Group)
498+
assert observed == subgroup
499+
assert group.get_group("subgroup/subsubgroup") == subsubgroup
500+
501+
502+
def test_group_get_group_missing(store: Store, zarr_format: ZarrFormat) -> None:
503+
"""
504+
`Group.get_group` raises `GroupNotFoundError` when no node exists at the given path.
505+
"""
506+
group = Group.from_store(store, zarr_format=zarr_format)
507+
with pytest.raises(GroupNotFoundError, match="No group found in store"):
508+
group.get_group("missing")
509+
510+
511+
def test_group_get_group_wrong_node_type(store: Store, zarr_format: ZarrFormat) -> None:
512+
"""
513+
`Group.get_group` raises `ContainsArrayError` when the node at the given path is an
514+
array rather than a group.
515+
"""
516+
group = Group.from_store(store, zarr_format=zarr_format)
517+
group.create_array(name="subarray", shape=(10,), chunks=(10,), dtype="uint8")
518+
with pytest.raises(ContainsArrayError, match="An array exists in store"):
519+
group.get_group("subarray")
520+
521+
449522
@pytest.mark.parametrize("consolidated", [True, False])
450523
def test_group_delitem(store: Store, zarr_format: ZarrFormat, consolidated: bool) -> None:
451524
"""
@@ -1469,7 +1542,7 @@ async def test_group_getitem_consolidated(self, store: Store) -> None:
14691542

14701543
# On disk, we've consolidated all the metadata in the root zarr.json
14711544
group = await zarr.api.asynchronous.open(store=store)
1472-
rg0 = await group.getitem("g0")
1545+
rg0 = await group.get_group("g0")
14731546

14741547
expected = ConsolidatedMetadata(
14751548
metadata={
@@ -1490,10 +1563,10 @@ async def test_group_getitem_consolidated(self, store: Store) -> None:
14901563
)
14911564
assert rg0.metadata.consolidated_metadata == expected
14921565

1493-
rg1 = await rg0.getitem("g1")
1566+
rg1 = await rg0.get_group("g1")
14941567
assert rg1.metadata.consolidated_metadata == expected.metadata["g1"].consolidated_metadata
14951568

1496-
rg2 = await rg1.getitem("g2")
1569+
rg2 = await rg1.get_group("g2")
14971570
assert rg2.metadata.consolidated_metadata == ConsolidatedMetadata(metadata={})
14981571

14991572
async def test_group_delitem_consolidated(self, store: Store) -> None:

tests/test_metadata/test_consolidated.py

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -111,12 +111,10 @@ async def test_getitem_consolidated_empty_leaf_group(
111111
group = await zarr.api.asynchronous.open_consolidated(
112112
store=memory_store, zarr_format=zarr_format
113113
)
114-
raw = await group.getitem("raw")
115-
assert isinstance(raw, zarr.AsyncGroup)
114+
raw = await group.get_group("raw")
116115
assert raw.metadata.consolidated_metadata is not None
117116

118-
varm = await raw.getitem("varm")
119-
assert isinstance(varm, zarr.AsyncGroup)
117+
varm = await raw.get_group("varm")
120118
assert varm.metadata.consolidated_metadata == ConsolidatedMetadata(metadata={})
121119

122120
async def test_open_consolidated_false_raises(self) -> None:
@@ -770,8 +768,7 @@ async def test_absolute_path_for_subgroup(self, memory_store: zarr.storage.Memor
770768
await zarr.api.asynchronous.consolidate_metadata(memory_store)
771769

772770
group = await zarr.api.asynchronous.open_group(store=memory_store)
773-
subgroup = await group.getitem("/a")
774-
assert isinstance(subgroup, AsyncGroup)
771+
subgroup = await group.get_group("/a")
775772
members = [x async for x in subgroup.keys()] # noqa: SIM118
776773
assert members == ["b"]
777774

tests/test_store/test_zip.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,6 @@
2222
import zarr
2323
from zarr import create_array
2424
from zarr.core.buffer import Buffer, cpu, default_buffer_prototype
25-
from zarr.core.group import Group
2625
from zarr.core.sync import sync
2726
from zarr.storage import ZipStore
2827
from zarr.testing.store import StoreTests
@@ -141,13 +140,13 @@ def test_externally_zipped_store(self, tmp_path: Path) -> None:
141140
zarr_path = tmp_path / "foo.zarr"
142141
root = zarr.open_group(store=zarr_path, mode="w")
143142
root.require_group("foo")
144-
assert isinstance(foo := root["foo"], Group) # noqa: RUF018
143+
foo = root.get_group("foo")
145144
foo["bar"] = np.array([1])
146145
shutil.make_archive(str(zarr_path), "zip", zarr_path)
147146
zip_path = tmp_path / "foo.zarr.zip"
148147
zipped = zarr.open_group(ZipStore(zip_path, mode="r"), mode="r")
149148
assert list(zipped.keys()) == list(root.keys())
150-
assert isinstance(group := zipped["foo"], Group)
149+
group = zipped.get_group("foo")
151150
assert list(group.keys()) == list(group.keys())
152151

153152
async def test_list_without_explicit_open(self, tmp_path: Path) -> None:

0 commit comments

Comments
 (0)