group-level create_hierarchy

Author: d-v-b

src/zarr/core/group.py

@@ -9,6 +9,7 @@
 from dataclasses import asdict, dataclass, field, fields, replace
 from functools import partial
 from itertools import accumulate
+from pathlib import PurePosixPath
 from typing import TYPE_CHECKING, Literal, TypeVar, assert_never, cast, overload
 
 import numpy as np
@@ -1424,6 +1425,33 @@ async def _members(
         ):
             yield member
 
+    # TODO: find a better name for this. create_tree could work.
+    # TODO: include an example in the docstring
+    async def create_hierarchy(
+        self, nodes: dict[str, ArrayV2Metadata | ArrayV3Metadata | GroupMetadata]
+    ) -> AsyncIterator[AsyncGroup | AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]]:
+        """
+        Create a hierarchy of arrays or groups rooted at this group.
+
+        This method takes a dictionary where the keys are the names of the arrays or groups
+        to create and the values are the metadata or objects representing the arrays or groups.
+
+        The method returns an asynchronous iterator over the created nodes.
+
+        Parameters
+        ----------
+        nodes : A dictionary representing the hierarchy to create
+
+        Returns
+        -------
+            An asynchronous iterator over the created nodes.
+        """
+        semaphore = asyncio.Semaphore(config.get("async.concurrency"))
+        async for node in create_hierarchy(
+            store_path=self.store_path, nodes=nodes, semaphore=semaphore
+        ):
+            yield node
+
     async def keys(self) -> AsyncGenerator[str, None]:
         """Iterate over member names."""
         async for key, _ in self.members():
@@ -2046,6 +2074,32 @@ def members(self, max_depth: int | None = 0) -> tuple[tuple[str, Array | Group],
 
         return tuple((kv[0], _parse_async_node(kv[1])) for kv in _members)
 
+    def create_hierarchy(
+        self, nodes: dict[str, ArrayV2Metadata | ArrayV3Metadata | GroupMetadata]
+    ) -> dict[str, AsyncGroup | AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]]:
+        """
+        Create a hierarchy of arrays or groups rooted at this group.
+
+        This method takes a dictionary where the keys are the names of the arrays or groups
+        to create and the values are the metadata objects for the arrays or groups.
+
+        The method returns an asynchronous iterator over the created nodes.
+
+        Parameters
+        ----------
+        nodes : A dictionary representing the hierarchy to create
+
+        Returns
+        -------
+            A dict containing the created nodes.The keys are the same as th
+        """
+        nodes_created = self._sync_iter(self._async_group.create_hierarchy(nodes))
+        if self.path == "":
+            root = "/"
+        else:
+            root = self.path
+        return {str(PurePosixPath(n.name).relative_to(root)): n for n in nodes_created}
+
     def keys(self) -> Generator[str, None]:
         """Return an iterator over group member names.
 

tests/test_group.py

@@ -1488,6 +1488,24 @@ async def test_create_hierarchy(store: Store, zarr_format: ZarrFormat) -> None:
     assert expected_meta == {k: v.metadata for k, v in observed_nodes.items()}
 
 
+@pytest.mark.parametrize("store", ["memory"], indirect=True)
+def test_group_create_hierarchy(store: Store, zarr_format: ZarrFormat):
+    """
+    Test that the Group.create_hierarchy method creates specified nodes and returns them in a dict.
+    """
+    g = Group.from_store(store)
+    tree = {
+        "a": GroupMetadata(zarr_format=zarr_format, attributes={"name": "a"}),
+        "a/b": GroupMetadata(zarr_format=zarr_format, attributes={"name": "a/b"}),
+        "a/b/c": meta_from_array(
+            np.zeros(5), zarr_format=zarr_format, attributes={"name": "a/b/c"}
+        ),
+    }
+    nodes = g.create_hierarchy(tree)
+    for k, v in nodes.items():
+        assert v.metadata == tree[k]
+
+
 def test_group_members_performance(store: MemoryStore) -> None:
     """
     Test that the execution time of Group.members is less than the number of members times the