Added Group.tree method

Author: TomAugspurger

docs/tutorial.rst

@@ -330,10 +330,10 @@ representation of the hierarchy, e.g.::
 
     >>> root.tree()
     /
-     └── foo
-         └── bar
-             ├── baz (10000, 10000) int32
-             └── quux (10000, 10000) int32
+    └── foo
+        └── bar
+            ├── baz (10000, 10000) int32
+            └── quux (10000, 10000) int32
 
 The :func:`zarr.convenience.open` function provides a convenient way to create or
 re-open a group stored in a directory on the file-system, with sub-groups stored in
@@ -424,6 +424,12 @@ Groups also have the :func:`zarr.hierarchy.Group.tree` method, e.g.::
          ├── bar (1000000,) int64
          └── baz (1000, 1000) float32
 
+
+.. note::
+
+   :func:`zarr.Group.tree` requires the optional `rich <https://rich.readthedocs.io/en/stable/>`_
+   dependency. It can be installed with the ``[tree]`` extra.
+
 If you're using Zarr within a Jupyter notebook (requires
 `ipytree <https://github.com/QuantStack/ipytree>`_), calling ``tree()`` will generate an
 interactive tree representation, see the `repr_tree.ipynb notebook

pyproject.toml

@@ -101,6 +101,9 @@ optional = [
     'lmdb',
     'universal-pathlib>=0.0.22',
 ]
+tree = [
+    "rich",
+]
 
 [project.urls]
 "Bug Tracker" = "https://github.com/zarr-developers/zarr-python/issues"

src/zarr/_tree.py

@@ -0,0 +1,44 @@
+import io
+
+from zarr.core.group import AsyncGroup
+
+try:
+    import rich
+    import rich.console
+    import rich.tree
+except ImportError as e:
+    raise ImportError("'rich' is required for Group.tree") from e
+
+
+class TreeRepr:
+    def __init__(self, tree: rich.tree.Tree) -> None:
+        self.tree = tree
+
+    def __repr__(self) -> str:
+        console = rich.console.Console(file=io.StringIO())
+        console.print(self.tree)
+        return str(console.file.getvalue())
+
+
+async def group_tree_async(group: AsyncGroup, max_depth: int | None = None) -> TreeRepr:
+    tree = rich.tree.Tree(label=f"[b]{group.name}[/b]")
+    nodes = {"": tree}
+    members = sorted([x async for x in group.members(max_depth=max_depth)])
+
+    for key, node in members:
+        if key.count("/") == 0:
+            parent_key = ""
+        else:
+            parent_key = key.rsplit("/", 1)[0]
+        parent = nodes[parent_key]
+
+        # We want what the spec calls the node "name", the part excluding all leading
+        # /'s and path segments. But node.name includes all that, so we build it here.
+        name = key.rsplit("/")[-1]
+        if isinstance(node, AsyncGroup):
+            label = f"[b]{name}[/b]"
+        else:
+            label = f"[b]{name}[/b] {node.shape} {node.dtype}"
+        nodes[key] = parent.add(label)
+
+    return TreeRepr(tree)

src/zarr/core/group.py

@@ -1270,8 +1270,30 @@ async def array_values(
         async for _, array in self.arrays():
             yield array
 
-    async def tree(self, expand: bool = False, level: int | None = None) -> Any:
-        raise NotImplementedError
+    async def tree(self, expand: bool | None = None, level: int | None = None) -> Any:
+        """
+        Return a tree-like representation of a hierarchy.
+
+        This requires the optional ``rich`` dependency.
+
+        Parameters
+        ----------
+        expand : bool, optional
+            This keyword is not yet supported. A NotImplementedError is raised if
+            it's used.
+        level : int, optional
+            The maximum depth below this Group to display in the tree.
+
+        Returns
+        -------
+        TreeRepr
+            A pretty-printable object displaying the hierarchy.
+        """
+        from zarr._tree import group_tree_async
+
+        if expand is not None:
+            raise NotImplementedError("'expanded' is not yet implemented.")
+        return await group_tree_async(self, max_depth=level)
 
     async def empty(
         self, *, name: str, shape: ChunkCoords, **kwargs: Any
@@ -1504,7 +1526,25 @@ def array_values(self) -> Generator[Array, None]:
         for _, array in self.arrays():
             yield array
 
-    def tree(self, expand: bool = False, level: int | None = None) -> Any:
+    def tree(self, expand: bool | None = None, level: int | None = None) -> Any:
+        """
+        Return a tree-like representation of a hierarchy.
+
+        This requires the optional ``rich`` dependency.
+
+        Parameters
+        ----------
+        expand : bool, optional
+            This keyword is not yet supported. A NotImplementedError is raised if
+            it's used.
+        level : int, optional
+            The maximum depth below this Group to display in the tree.
+
+        Returns
+        -------
+        TreeRepr
+            A pretty-printable object displaying the hierarchy.
+        """
         return self._sync(self._async_group.tree(expand=expand, level=level))
 
     def create_group(self, name: str, **kwargs: Any) -> Group:

tests/test_tree.py

@@ -0,0 +1,54 @@
+import textwrap
+from typing import Any
+
+import pytest
+
+import zarr
+
+
+@pytest.mark.parametrize("root_name", [None, "root"])
+def test_tree(root_name: Any) -> None:
+    g = zarr.group(path=root_name)
+    A = g.create_group("A")
+    B = g.create_group("B")
+    C = B.create_group("C")
+    D = C.create_group("C")
+
+    A.create_array(name="x", shape=(2), dtype="float64")
+    A.create_array(name="y", shape=(0,), dtype="int8")
+    B.create_array(name="x", shape=(0,))
+    C.create_array(name="x", shape=(0,))
+    D.create_array(name="x", shape=(0,))
+
+    result = repr(g.tree())
+    root = root_name or ""
+
+    expected = textwrap.dedent(f"""\
+        /{root}
+        ├── A
+        │   ├── x (2,) float64
+        │   └── y (0,) int8
+        └── B
+            ├── C
+            │   ├── C
+            │   │   └── x (0,) float64
+            │   └── x (0,) float64
+            └── x (0,) float64
+        """)
+
+    assert result == expected
+
+    result = repr(g.tree(level=0))
+    expected = textwrap.dedent(f"""\
+        /{root}
+        ├── A
+        └── B
+        """)
+
+    assert result == expected
+
+
+def test_expand_not_implemented() -> None:
+    g = zarr.group()
+    with pytest.raises(NotImplementedError):
+        g.tree(expand=True)