diff --git a/changes/4128.feature.md b/changes/4128.feature.md new file mode 100644 index 0000000000..c62a615ac2 --- /dev/null +++ b/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. diff --git a/docs/user-guide/groups.md b/docs/user-guide/groups.md index 5faa26a281..40199be94f 100644 --- a/docs/user-guide/groups.md +++ b/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.: diff --git a/src/zarr/core/group.py b/src/zarr/core/group.py index 0aaf89234e..922eaf1498 100644 --- a/src/zarr/core/group.py +++ b/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") + # + ``` + """ + 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") + # + ``` + """ + return Group(self._sync(self._async_group.get_group(path))) + def __delitem__(self, key: str) -> None: """Delete a group member. diff --git a/tests/test_group.py b/tests/test_group.py index bc80e19e86..5a14fa6dc0 100644 --- a/tests/test_group.py +++ b/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: diff --git a/tests/test_metadata/test_consolidated.py b/tests/test_metadata/test_consolidated.py index 3596d2bcaa..e6087435fe 100644 --- a/tests/test_metadata/test_consolidated.py +++ b/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"] diff --git a/tests/test_store/test_zip.py b/tests/test_store/test_zip.py index ed69114b51..740877c792 100644 --- a/tests/test_store/test_zip.py +++ 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: