From 8c412c118dc78f3a069611f68f0d34f44a7694e4 Mon Sep 17 00:00:00 2001 From: Ian Hunt-Isaak Date: Fri, 6 Mar 2026 11:45:57 -0500 Subject: [PATCH 1/3] remove rich dependency for tree --- changes/XXXX.feature.md | 1 + docs/user-guide/groups.md | 2 - docs/user-guide/installation.md | 2 +- pyproject.toml | 4 +- src/zarr/__init__.py | 1 - src/zarr/core/_tree.py | 126 ++++++++++++++++++++++++-------- src/zarr/core/group.py | 38 ++++++++-- tests/test_api.py | 1 - tests/test_tree.py | 64 +++++++++++++--- 9 files changed, 181 insertions(+), 58 deletions(-) create mode 100644 changes/XXXX.feature.md diff --git a/changes/XXXX.feature.md b/changes/XXXX.feature.md new file mode 100644 index 0000000000..17f26666ed --- /dev/null +++ b/changes/XXXX.feature.md @@ -0,0 +1 @@ +`Group.tree()` no longer requires the `rich` dependency. Tree rendering now uses built-in ANSI bold for terminals and HTML bold for Jupyter. New parameters: `plain=True` for unstyled output, and `max_nodes` (default 500) to truncate large hierarchies with early bailout. diff --git a/docs/user-guide/groups.md b/docs/user-guide/groups.md index e093590dfe..58a9c1c806 100644 --- a/docs/user-guide/groups.md +++ b/docs/user-guide/groups.md @@ -133,5 +133,3 @@ Groups also have the [`zarr.Group.tree`][] method, e.g.: print(root.tree()) ``` -!!! note - [`zarr.Group.tree`][] requires the optional [rich](https://rich.readthedocs.io/en/stable/) dependency. It can be installed with the `[tree]` extra. \ No newline at end of file diff --git a/docs/user-guide/installation.md b/docs/user-guide/installation.md index 4d323643f1..6c1414e81a 100644 --- a/docs/user-guide/installation.md +++ b/docs/user-guide/installation.md @@ -26,7 +26,7 @@ These can be installed using `pip install "zarr[]"`, e.g. `pip install "z - `gpu`: support for GPUs - `remote`: support for reading/writing to remote data stores -Additional optional dependencies include `rich`, `universal_pathlib`. These must be installed separately. +Additional optional dependencies include `universal_pathlib`. These must be installed separately. ## conda diff --git a/pyproject.toml b/pyproject.toml index 18bdeda07c..4853636ef9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -70,7 +70,7 @@ gpu = [ "cupy-cuda12x", ] cli = ["typer"] -optional = ["rich", "universal-pathlib"] +optional = ["universal-pathlib"] [project.scripts] zarr = "zarr._cli.cli:app" @@ -122,7 +122,6 @@ docs = [ "towncrier", # Optional dependencies to run examples "numcodecs[msgpack]", - "rich", "s3fs>=2023.10.0", "astroid<4", "pytest", @@ -131,7 +130,6 @@ dev = [ {include-group = "test"}, {include-group = "remote-tests"}, {include-group = "docs"}, - "rich", "universal-pathlib", "mypy", ] diff --git a/src/zarr/__init__.py b/src/zarr/__init__.py index 3c6195c28f..cdf3840c3b 100644 --- a/src/zarr/__init__.py +++ b/src/zarr/__init__.py @@ -78,7 +78,6 @@ def print_packages(packages: list[str]) -> None: "s3fs", "gcsfs", "universal-pathlib", - "rich", "obstore", ] diff --git a/src/zarr/core/_tree.py b/src/zarr/core/_tree.py index eed807ec95..fdf69404ad 100644 --- a/src/zarr/core/_tree.py +++ b/src/zarr/core/_tree.py @@ -1,17 +1,11 @@ -import io -import os +import sys +from collections import deque from collections.abc import Sequence +from html import escape as html_escape from typing import Any 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: """ @@ -21,45 +15,115 @@ class TreeRepr: of Zarr's public API. """ - def __init__(self, tree: rich.tree.Tree) -> None: - self._tree = tree + def __init__(self, text: str, html: str) -> None: + self._text = text + self._html = html def __repr__(self) -> str: - color_system = os.environ.get("OVERRIDE_COLOR_SYSTEM", rich.get_console().color_system) - console = rich.console.Console(file=io.StringIO(), color_system=color_system) - console.print(self._tree) - return str(console.file.getvalue()) + return self._text def _repr_mimebundle_( self, - include: Sequence[str], - exclude: Sequence[str], + include: Sequence[str] | None = None, + exclude: Sequence[str] | None = None, **kwargs: Any, ) -> dict[str, str]: # For jupyter support. - # Unsure why mypy infers the return type to by Any - return self._tree._repr_mimebundle_(include=include, exclude=exclude, **kwargs) # type: ignore[no-any-return] + html = ( + '
"
+            f"{self._html}
\n" + ) + return {"text/plain": self._text, "text/html": html} + +async def group_tree_async( + group: AsyncGroup, + max_depth: int | None = None, + max_nodes: int = 500, + plain: bool = False, +) -> TreeRepr: + members = [] + truncated = False + async for item in group.members(max_depth=max_depth): + if len(members) == max_nodes: + truncated = True + break + members.append(item) + members.sort(key=lambda key_node: key_node[0]) -async def group_tree_async(group: AsyncGroup, max_depth: int | None = None) -> TreeRepr: - tree = rich.tree.Tree(label=f"[bold]{group.name}[/bold]") - nodes = {"": tree} - members = sorted([x async for x in group.members(max_depth=max_depth)]) + # Set up styling tokens: ANSI bold for terminals, HTML for Jupyter, + # or empty strings when plain=True (useful for LLMs, logging, files). + if plain: + ansi_open = ansi_close = html_open = html_close = "" + else: + # Avoid emitting ANSI escape codes when output is piped or in CI. + use_ansi = sys.stdout.isatty() + ansi_open = "\x1b[1m" if use_ansi else "" + ansi_close = "\x1b[0m" if use_ansi else "" + html_open = "" + html_close = "" + if truncated: + note = f"Truncated at max_nodes={max_nodes}, some nodes and their children may be missing\n" + text_lines = [note, f"{ansi_open}{group.name}{ansi_close}"] + html_lines = [note, f"{html_open}{html_escape(group.name)}{html_close}"] + else: + text_lines = [f"{ansi_open}{group.name}{ansi_close}"] + html_lines = [f"{html_open}{html_escape(group.name)}{html_close}"] + + # Group members by parent key so we can render the tree level by level. + nodes: dict[str, list[tuple[str, Any]]] = {} for key, node in members: if key.count("/") == 0: parent_key = "" else: parent_key = key.rsplit("/", 1)[0] - parent = nodes[parent_key] + nodes.setdefault(parent_key, []).append((key, node)) - # 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. + # Render the tree iteratively (not recursively) to avoid hitting + # Python's recursion limit on deeply nested hierarchies. + # Each stack frame is (prefix_string, remaining_children_at_this_level). + stack = [("", deque(nodes.get("", [])))] + while stack: + prefix, remaining = stack[-1] + if not remaining: + stack.pop() + continue + key, node = remaining.popleft() name = key.rsplit("/")[-1] + escaped_name = html_escape(name) + # if we popped the last item then remaining will + # now be empty - that's how we got past the if not remaining + # above, but this can still be true. + is_last = not remaining + connector = "└── " if is_last else "├── " if isinstance(node, AsyncGroup): - label = f"[bold]{name}[/bold]" + text_lines.append(f"{prefix}{connector}{ansi_open}{name}{ansi_close}") + html_lines.append(f"{prefix}{connector}{html_open}{escaped_name}{html_close}") else: - label = f"[bold]{name}[/bold] {node.shape} {node.dtype}" - nodes[key] = parent.add(label) - - return TreeRepr(tree) + text_lines.append( + f"{prefix}{connector}{ansi_open}{name}{ansi_close} {node.shape} {node.dtype}" + ) + html_lines.append( + f"{prefix}{connector}{html_open}{escaped_name}{html_close}" + f" {html_escape(str(node.shape))} {html_escape(str(node.dtype))}" + ) + # Descend into children with an accumulated prefix: + # Example showing how prefix accumulates: + # / + # ├── a prefix = "" + # │ ├── b prefix = "" + "│ " + # │ │ └── x prefix = "" + "│ " + "│ " + # │ └── c prefix = "" + "│ " + # └── d prefix = "" + # └── e prefix = "" + " " + if children := nodes.get(key, []): + if is_last: + child_prefix = prefix + " " + else: + child_prefix = prefix + "│ " + stack.append((child_prefix, deque(children))) + text = "\n".join(text_lines) + "\n" + html = "\n".join(html_lines) + "\n" + return TreeRepr(text, html) diff --git a/src/zarr/core/group.py b/src/zarr/core/group.py index 9b5fee275b..841500275e 100644 --- a/src/zarr/core/group.py +++ b/src/zarr/core/group.py @@ -1588,12 +1588,16 @@ async def array_values( async for _, array in self.arrays(): yield array - async def tree(self, expand: bool | None = None, level: int | None = None) -> Any: + async def tree( + self, + expand: bool | None = None, + level: int | None = None, + max_nodes: int = 500, + plain: bool = False, + ) -> Any: """ Return a tree-like representation of a hierarchy. - This requires the optional ``rich`` dependency. - Parameters ---------- expand : bool, optional @@ -1601,6 +1605,12 @@ async def tree(self, expand: bool | None = None, level: int | None = None) -> An it's used. level : int, optional The maximum depth below this Group to display in the tree. + max_nodes : int + Maximum number of nodes to display before truncating. Default is 500. + plain : bool, optional + If True, return a plain-text tree without ANSI styling. This is + useful when the output will be consumed by an LLM or written to a + file. Default is False. Returns ------- @@ -1611,7 +1621,7 @@ async def tree(self, expand: bool | None = None, level: int | None = None) -> An if expand is not None: raise NotImplementedError("'expand' is not yet implemented.") - return await group_tree_async(self, max_depth=level) + return await group_tree_async(self, max_depth=level, max_nodes=max_nodes, plain=plain) async def empty(self, *, name: str, shape: tuple[int, ...], **kwargs: Any) -> AnyAsyncArray: """Create an empty array with the specified shape in this Group. The contents will @@ -2371,12 +2381,16 @@ def array_values(self) -> Generator[AnyArray, None]: for _, array in self.arrays(): yield array - def tree(self, expand: bool | None = None, level: int | None = None) -> Any: + def tree( + self, + expand: bool | None = None, + level: int | None = None, + max_nodes: int = 500, + plain: bool = False, + ) -> Any: """ Return a tree-like representation of a hierarchy. - This requires the optional ``rich`` dependency. - Parameters ---------- expand : bool, optional @@ -2384,13 +2398,21 @@ def tree(self, expand: bool | None = None, level: int | None = None) -> Any: it's used. level : int, optional The maximum depth below this Group to display in the tree. + max_nodes : int + Maximum number of nodes to display before truncating. Default is 500. + plain : bool, optional + If True, return a plain-text tree without ANSI styling. This is + useful when the output will be consumed by an LLM or written to a + file. Default is False. Returns ------- TreeRepr A pretty-printable object displaying the hierarchy. """ - return self._sync(self._async_group.tree(expand=expand, level=level)) + return self._sync( + self._async_group.tree(expand=expand, level=level, max_nodes=max_nodes, plain=plain) + ) def create_group(self, name: str, **kwargs: Any) -> Group: """Create a sub-group. diff --git a/tests/test_api.py b/tests/test_api.py index 07c3c8590d..a306ff3dc3 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -599,7 +599,6 @@ def test_load_local(tmp_path: Path, path: str | None, load_read_only: bool) -> N def test_tree() -> None: - pytest.importorskip("rich") g1 = zarr.group() g1.create_group("foo") g3 = g1.create_group("bar") diff --git a/tests/test_tree.py b/tests/test_tree.py index b4a5106998..78ea121f4d 100644 --- a/tests/test_tree.py +++ b/tests/test_tree.py @@ -1,4 +1,3 @@ -import os import textwrap from typing import Any @@ -6,12 +5,19 @@ import zarr -pytest.importorskip("rich") - @pytest.mark.parametrize("root_name", [None, "root"]) -def test_tree(root_name: Any) -> None: - os.environ["OVERRIDE_COLOR_SYSTEM"] = "truecolor" +@pytest.mark.parametrize("atty", [True, False]) +@pytest.mark.parametrize("plain", [True, False]) +def test_tree(root_name: Any, atty: bool, plain: bool, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("sys.stdout.isatty", lambda: atty) + + if atty and not plain: + BOPEN = "\x1b[1m" + BCLOSE = "\x1b[0m" + else: + BOPEN = "" + BCLOSE = "" g = zarr.group(path=root_name) A = g.create_group("A") @@ -25,12 +31,9 @@ def test_tree(root_name: Any) -> None: C.create_array(name="x", shape=(0,), dtype="float64") D.create_array(name="x", shape=(0,), dtype="float64") - result = repr(g.tree()) + result = repr(g.tree(plain=plain)) root = root_name or "" - BOPEN = "\x1b[1m" - BCLOSE = "\x1b[0m" - expected = textwrap.dedent(f"""\ {BOPEN}/{root}{BCLOSE} ├── {BOPEN}A{BCLOSE} @@ -46,15 +49,54 @@ def test_tree(root_name: Any) -> None: assert result == expected - result = repr(g.tree(level=0)) + result = repr(g.tree(level=0, plain=plain)) expected = textwrap.dedent(f"""\ {BOPEN}/{root}{BCLOSE} ├── {BOPEN}A{BCLOSE} └── {BOPEN}B{BCLOSE} """) - assert result == expected + if not plain: + tree = g.tree(plain=False) + bundle = tree._repr_mimebundle_() + assert "text/plain" in bundle + assert "text/html" in bundle + assert "A" in bundle["text/html"] + assert "x" in bundle["text/html"] + assert " None: + g = zarr.group() + g.create_group("a") + g.create_group("b") + g.create_group("c") + g.create_group("d") + g.create_group("e") + + result = repr(g.tree(max_nodes=3, plain=True)) + assert "Truncated at max_nodes=3" in result + # Should show exactly 3 nodes (lines with ── connectors). + lines = result.strip().split("\n") + node_lines = [line for line in lines if "──" in line] + assert len(node_lines) == 3 + + # Full tree should not show truncation message. + full = repr(g.tree(max_nodes=500, plain=True)) + assert "truncated" not in full + + +def test_tree_html_escaping() -> None: + g = zarr.group() + g.create_group("") + + tree = g.tree() + bundle = tree._repr_mimebundle_() + assert "<img" in bundle["text/html"] + assert "" in bundle["text/plain"] + def test_expand_not_implemented() -> None: g = zarr.group() From 8cdea42c06b66381dc29a1600275214480d60f91 Mon Sep 17 00:00:00 2001 From: Ian Hunt-Isaak Date: Fri, 13 Mar 2026 12:01:47 -0400 Subject: [PATCH 2/3] better --- src/zarr/core/_tree.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/src/zarr/core/_tree.py b/src/zarr/core/_tree.py index fdf69404ad..215e2097af 100644 --- a/src/zarr/core/_tree.py +++ b/src/zarr/core/_tree.py @@ -15,11 +15,14 @@ class TreeRepr: of Zarr's public API. """ - def __init__(self, text: str, html: str) -> None: + def __init__(self, text: str, html: str, truncated: str = "") -> None: self._text = text self._html = html + self._truncated = truncated def __repr__(self) -> str: + if self._truncated: + return self._truncated + self._text return self._text def _repr_mimebundle_( @@ -28,13 +31,15 @@ def _repr_mimebundle_( exclude: Sequence[str] | None = None, **kwargs: Any, ) -> dict[str, str]: + text = self._truncated + self._text if self._truncated else self._text # For jupyter support. + html_body = self._truncated + self._html if self._truncated else self._html html = ( '
"
-            f"{self._html}
\n" + f"{html_body}\n" ) - return {"text/plain": self._text, "text/html": html} + return {"text/plain": text, "text/html": html} async def group_tree_async( @@ -64,14 +69,6 @@ async def group_tree_async( html_open = "" html_close = "" - if truncated: - note = f"Truncated at max_nodes={max_nodes}, some nodes and their children may be missing\n" - text_lines = [note, f"{ansi_open}{group.name}{ansi_close}"] - html_lines = [note, f"{html_open}{html_escape(group.name)}{html_close}"] - else: - text_lines = [f"{ansi_open}{group.name}{ansi_close}"] - html_lines = [f"{html_open}{html_escape(group.name)}{html_close}"] - # Group members by parent key so we can render the tree level by level. nodes: dict[str, list[tuple[str, Any]]] = {} for key, node in members: @@ -84,6 +81,8 @@ async def group_tree_async( # Render the tree iteratively (not recursively) to avoid hitting # Python's recursion limit on deeply nested hierarchies. # Each stack frame is (prefix_string, remaining_children_at_this_level). + text_lines = [f"{ansi_open}{group.name}{ansi_close}"] + html_lines = [f"{html_open}{html_escape(group.name)}{html_close}"] stack = [("", deque(nodes.get("", [])))] while stack: prefix, remaining = stack[-1] @@ -126,4 +125,5 @@ async def group_tree_async( stack.append((child_prefix, deque(children))) text = "\n".join(text_lines) + "\n" html = "\n".join(html_lines) + "\n" - return TreeRepr(text, html) + note = f"Truncated at max_nodes={max_nodes}, some nodes and their children may be missing\n" if truncated else "" + return TreeRepr(text, html, truncated=note) From 8aeebaa426a817ba2936b10fe1e5d52c33a3c1ce Mon Sep 17 00:00:00 2001 From: Ian Hunt-Isaak Date: Fri, 13 Mar 2026 14:13:10 -0400 Subject: [PATCH 3/3] lint + pr number --- changes/{XXXX.feature.md => 3778.misc.md} | 0 src/zarr/core/_tree.py | 8 ++++++-- 2 files changed, 6 insertions(+), 2 deletions(-) rename changes/{XXXX.feature.md => 3778.misc.md} (100%) diff --git a/changes/XXXX.feature.md b/changes/3778.misc.md similarity index 100% rename from changes/XXXX.feature.md rename to changes/3778.misc.md diff --git a/src/zarr/core/_tree.py b/src/zarr/core/_tree.py index 215e2097af..61cfc57ecc 100644 --- a/src/zarr/core/_tree.py +++ b/src/zarr/core/_tree.py @@ -48,7 +48,7 @@ async def group_tree_async( max_nodes: int = 500, plain: bool = False, ) -> TreeRepr: - members = [] + members: list[tuple[str, Any]] = [] truncated = False async for item in group.members(max_depth=max_depth): if len(members) == max_nodes: @@ -125,5 +125,9 @@ async def group_tree_async( stack.append((child_prefix, deque(children))) text = "\n".join(text_lines) + "\n" html = "\n".join(html_lines) + "\n" - note = f"Truncated at max_nodes={max_nodes}, some nodes and their children may be missing\n" if truncated else "" + note = ( + f"Truncated at max_nodes={max_nodes}, some nodes and their children may be missing\n" + if truncated + else "" + ) return TreeRepr(text, html, truncated=note)