feat: add type-safe get_array and get_group methods to AsyncGroup and Group

changes/4128.feature.md

@@ -0,0 +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

@@ -43,6 +43,29 @@ print(root['foo/bar'])
 print(root['foo/bar/baz'])
 ```
 
+Accessing a member with `[]` returns either an [`zarr.Array`][] or a [`zarr.Group`][], depending on
+what is stored at the given path. When you expect a node of a particular kind, use
+[`zarr.Group.get_array`][] or [`zarr.Group.get_group`][] instead. These methods accept the same
+paths as `[]`, but they have precise return types and raise an error if no node exists at the
+given path, or if the node is not of the expected kind:
+
+```python exec="true" session="groups" source="above" result="ansi"
+print(root.get_group('foo'))
+```
+
+```python exec="true" session="groups" source="above" result="ansi"
+print(root.get_array('foo/bar/baz'))
+```
+
+```python exec="true" session="groups" source="above" result="ansi"
+from zarr.errors import ContainsGroupError
+
+try:
+    root.get_array('foo')
+except ContainsGroupError as e:
+    print(e)
+```
+
 The [`zarr.Group.tree`][] method can be used to print a tree
 representation of the hierarchy, e.g.:
 

src/zarr/core/group.py

@@ -51,6 +51,7 @@
 from zarr.core.metadata.io import save_metadata
 from zarr.core.sync import SyncMixin, sync
 from zarr.errors import (
+    ArrayNotFoundError,
     ContainsArrayError,
     ContainsGroupError,
     GroupNotFoundError,
@@ -820,6 +821,70 @@ async def get[DefaultT](
         except KeyError:
             return default
 
+    async def get_array(self, path: str) -> AnyAsyncArray:
+        """Obtain an array member of this group, raising if it is absent or not an array.
+
+        Parameters
+        ----------
+        path : str
+            Path of the array relative to this group. May contain `/` to reference
+            a member of a subgroup, e.g. `subgroup/subarray`.
+
+        Returns
+        -------
+        AsyncArray
+            The array at the given path.
+
+        Raises
+        ------
+        ArrayNotFoundError
+            If no node exists at the given path.
+        ContainsGroupError
+            If the node at the given path is a group rather than an array.
+        """
+        store_path = self.store_path / path
+        try:
+            node = await self.getitem(path)
+        except KeyError as e:
+            msg = f"No array found in store {store_path.store!r} at path {store_path.path!r}"
+            raise ArrayNotFoundError(msg) from e
+        if isinstance(node, AsyncGroup):
+            msg = f"A group exists in store {store_path.store!r} at path {store_path.path!r}."
+            raise ContainsGroupError(msg)
+        return node
+
+    async def get_group(self, path: str) -> AsyncGroup:
+        """Obtain a group member of this group, raising if it is absent or not a group.
+
+        Parameters
+        ----------
+        path : str
+            Path of the group relative to this group. May contain `/` to reference
+            a member of a subgroup, e.g. `subgroup/subsubgroup`.
+
+        Returns
+        -------
+        AsyncGroup
+            The group at the given path.
+
+        Raises
+        ------
+        GroupNotFoundError
+            If no node exists at the given path.
+        ContainsArrayError
+            If the node at the given path is an array rather than a group.
+        """
+        store_path = self.store_path / path
+        try:
+            node = await self.getitem(path)
+        except KeyError as e:
+            msg = f"No group found in store {store_path.store!r} at path {store_path.path!r}"
+            raise GroupNotFoundError(msg) from e
+        if isinstance(node, AsyncArray):
+            msg = f"An array exists in store {store_path.store!r} at path {store_path.path!r}."
+            raise ContainsArrayError(msg)
+        return node
+
     async def _save_metadata(self, ensure_parents: bool = False) -> None:
         await save_metadata(self.store_path, self.metadata, ensure_parents=ensure_parents)
 
@@ -1880,6 +1945,74 @@ def get[DefaultT](
         except KeyError:
             return default
 
+    def get_array(self, path: str) -> AnyArray:
+        """Obtain an array member of this group, raising if it is absent or not an array.
+
+        Parameters
+        ----------
+        path : str
+            Path of the array relative to this group. May contain `/` to reference
+            a member of a subgroup, e.g. `subgroup/subarray`.
+
+        Returns
+        -------
+        Array
+            The array at the given path.
+
+        Raises
+        ------
+        ArrayNotFoundError
+            If no node exists at the given path.
+        ContainsGroupError
+            If the node at the given path is a group rather than an array.
+
+        Examples
+        --------
+        ```python
+        import zarr
+        from zarr.core.group import Group
+        group = Group.from_store(zarr.storage.MemoryStore())
+        group.create_array(name="subarray", shape=(10,), chunks=(10,), dtype="float64")
+        group.get_array("subarray")
+        # <Array memory://... shape=(10,) dtype=float64>
+        ```
+        """
+        return Array(self._sync(self._async_group.get_array(path)))
+
+    def get_group(self, path: str) -> Group:
+        """Obtain a group member of this group, raising if it is absent or not a group.
+
+        Parameters
+        ----------
+        path : str
+            Path of the group relative to this group. May contain `/` to reference
+            a member of a subgroup, e.g. `subgroup/subsubgroup`.
+
+        Returns
+        -------
+        Group
+            The group at the given path.
+
+        Raises
+        ------
+        GroupNotFoundError
+            If no node exists at the given path.
+        ContainsArrayError
+            If the node at the given path is an array rather than a group.
+
+        Examples
+        --------
+        ```python
+        import zarr
+        from zarr.core.group import Group
+        group = Group.from_store(zarr.storage.MemoryStore())
+        group.create_group(name="subgroup")
+        group.get_group("subgroup")
+        # <Group memory://...>
+        ```
+        """
+        return Group(self._sync(self._async_group.get_group(path)))
+
     def __delitem__(self, key: str) -> None:
         """Delete a group member.
 

tests/test_group.py

@@ -41,8 +41,10 @@
 from zarr.core.metadata.v3 import ArrayV3Metadata
 from zarr.core.sync import _collect_aiterator, sync
 from zarr.errors import (
+    ArrayNotFoundError,
     ContainsArrayError,
     ContainsGroupError,
+    GroupNotFoundError,
     MetadataValidationError,
     ZarrUserWarning,
 )
@@ -101,7 +103,7 @@ async def test_create_creates_parents(store: Store, zarr_format: ZarrFormat) ->
     root = await zarr.api.asynchronous.open_group(
         store=store,
     )
-    agroup = await root.getitem("a")
+    agroup = await root.get_group("a")
     assert agroup.attrs == {"key": "value"}
 
     # create a child node with a couple intermediates
@@ -446,6 +448,77 @@ def test_group_get_with_default(store: Store, zarr_format: ZarrFormat) -> None:
     assert result.attrs["foo"] == "bar"
 
 
+def test_group_get_array(store: Store, zarr_format: ZarrFormat) -> None:
+    """
+    `Group.get_array` returns the array at the given path, for both direct child names
+    and nested paths, and the result is statically typed as an Array.
+    """
+    group = Group.from_store(store, zarr_format=zarr_format)
+    subgroup = group.create_group(name="subgroup")
+    subarray = group.create_array(name="subarray", shape=(10,), chunks=(10,), dtype="uint8")
+    subsubarray = subgroup.create_array(name="subarray", shape=(10,), chunks=(10,), dtype="uint8")
+
+    observed = group.get_array("subarray")
+    assert isinstance(observed, Array)
+    assert observed == subarray
+    assert group.get_array("subgroup/subarray") == subsubarray
+
+
+def test_group_get_array_missing(store: Store, zarr_format: ZarrFormat) -> None:
+    """
+    `Group.get_array` raises `ArrayNotFoundError` when no node exists at the given path.
+    """
+    group = Group.from_store(store, zarr_format=zarr_format)
+    with pytest.raises(ArrayNotFoundError, match="No array found in store"):
+        group.get_array("missing")
+
+
+def test_group_get_array_wrong_node_type(store: Store, zarr_format: ZarrFormat) -> None:
+    """
+    `Group.get_array` raises `ContainsGroupError` when the node at the given path is a
+    group rather than an array.
+    """
+    group = Group.from_store(store, zarr_format=zarr_format)
+    group.create_group(name="subgroup")
+    with pytest.raises(ContainsGroupError, match="A group exists in store"):
+        group.get_array("subgroup")
+
+
+def test_group_get_group(store: Store, zarr_format: ZarrFormat) -> None:
+    """
+    `Group.get_group` returns the group at the given path, for both direct child names
+    and nested paths, and the result is statically typed as a Group.
+    """
+    group = Group.from_store(store, zarr_format=zarr_format)
+    subgroup = group.create_group(name="subgroup")
+    subsubgroup = subgroup.create_group(name="subsubgroup")
+
+    observed = group.get_group("subgroup")
+    assert isinstance(observed, Group)
+    assert observed == subgroup
+    assert group.get_group("subgroup/subsubgroup") == subsubgroup
+
+
+def test_group_get_group_missing(store: Store, zarr_format: ZarrFormat) -> None:
+    """
+    `Group.get_group` raises `GroupNotFoundError` when no node exists at the given path.
+    """
+    group = Group.from_store(store, zarr_format=zarr_format)
+    with pytest.raises(GroupNotFoundError, match="No group found in store"):
+        group.get_group("missing")
+
+
+def test_group_get_group_wrong_node_type(store: Store, zarr_format: ZarrFormat) -> None:
+    """
+    `Group.get_group` raises `ContainsArrayError` when the node at the given path is an
+    array rather than a group.
+    """
+    group = Group.from_store(store, zarr_format=zarr_format)
+    group.create_array(name="subarray", shape=(10,), chunks=(10,), dtype="uint8")
+    with pytest.raises(ContainsArrayError, match="An array exists in store"):
+        group.get_group("subarray")
+
+
 @pytest.mark.parametrize("consolidated", [True, False])
 def test_group_delitem(store: Store, zarr_format: ZarrFormat, consolidated: bool) -> None:
     """
@@ -1469,7 +1542,7 @@ async def test_group_getitem_consolidated(self, store: Store) -> None:
 
         # On disk, we've consolidated all the metadata in the root zarr.json
         group = await zarr.api.asynchronous.open(store=store)
-        rg0 = await group.getitem("g0")
+        rg0 = await group.get_group("g0")
 
         expected = ConsolidatedMetadata(
             metadata={
@@ -1490,10 +1563,10 @@ async def test_group_getitem_consolidated(self, store: Store) -> None:
         )
         assert rg0.metadata.consolidated_metadata == expected
 
-        rg1 = await rg0.getitem("g1")
+        rg1 = await rg0.get_group("g1")
         assert rg1.metadata.consolidated_metadata == expected.metadata["g1"].consolidated_metadata
 
-        rg2 = await rg1.getitem("g2")
+        rg2 = await rg1.get_group("g2")
         assert rg2.metadata.consolidated_metadata == ConsolidatedMetadata(metadata={})
 
     async def test_group_delitem_consolidated(self, store: Store) -> None:

tests/test_metadata/test_consolidated.py

@@ -111,12 +111,10 @@ async def test_getitem_consolidated_empty_leaf_group(
         group = await zarr.api.asynchronous.open_consolidated(
             store=memory_store, zarr_format=zarr_format
         )
-        raw = await group.getitem("raw")
-        assert isinstance(raw, zarr.AsyncGroup)
+        raw = await group.get_group("raw")
         assert raw.metadata.consolidated_metadata is not None
 
-        varm = await raw.getitem("varm")
-        assert isinstance(varm, zarr.AsyncGroup)
+        varm = await raw.get_group("varm")
         assert varm.metadata.consolidated_metadata == ConsolidatedMetadata(metadata={})
 
     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
             await zarr.api.asynchronous.consolidate_metadata(memory_store)
 
         group = await zarr.api.asynchronous.open_group(store=memory_store)
-        subgroup = await group.getitem("/a")
-        assert isinstance(subgroup, AsyncGroup)
+        subgroup = await group.get_group("/a")
         members = [x async for x in subgroup.keys()]  # noqa: SIM118
         assert members == ["b"]
 

tests/test_store/test_zip.py

@@ -20,7 +20,6 @@
 import zarr
 from zarr import create_array
 from zarr.core.buffer import Buffer, cpu, default_buffer_prototype
-from zarr.core.group import Group
 from zarr.core.sync import sync
 from zarr.storage import ZipStore
 from zarr.testing.store import StoreTests
@@ -139,13 +138,13 @@ def test_externally_zipped_store(self, tmp_path: Path) -> None:
         zarr_path = tmp_path / "foo.zarr"
         root = zarr.open_group(store=zarr_path, mode="w")
         root.require_group("foo")
-        assert isinstance(foo := root["foo"], Group)  # noqa: RUF018
+        foo = root.get_group("foo")
         foo["bar"] = np.array([1])
         shutil.make_archive(str(zarr_path), "zip", zarr_path)
         zip_path = tmp_path / "foo.zarr.zip"
         zipped = zarr.open_group(ZipStore(zip_path, mode="r"), mode="r")
         assert list(zipped.keys()) == list(root.keys())
-        assert isinstance(group := zipped["foo"], Group)
+        group = zipped.get_group("foo")
         assert list(group.keys()) == list(group.keys())
 
     async def test_list_without_explicit_open(self, tmp_path: Path) -> None: