fixup

Author: TomAugspurger

pyproject.toml

@@ -100,9 +100,10 @@ extra = [
 optional = [
     'lmdb',
     'universal-pathlib>=0.0.22',
+    'rich'
 ]
 tree = [
-    "rich",
+    'rich',
 ]
 
 [project.urls]

src/zarr/_tree.py

@@ -1,4 +1,5 @@
 import io
+from typing import Any
 
 from zarr.core.group import AsyncGroup
 
@@ -15,13 +16,19 @@ def __init__(self, tree: rich.tree.Tree) -> None:
         self.tree = tree
 
     def __repr__(self) -> str:
-        console = rich.console.Console(file=io.StringIO())
+        terminal = rich.get_console()
+        console = rich.console.Console(file=io.StringIO(), color_system=terminal.color_system)
         console.print(self.tree)
         return str(console.file.getvalue())
 
+    def _repr_mimebundle_(self, **kwargs: dict[str, Any]) -> dict[str, str]:
+        # For jupyter support.
+        # We don't depend on jupyter, so we can't do the static types appropriately here.
+        return self.tree._repr_mimebundle_(**kwargs)  # type: ignore[no-any-return]
+
 
 async def group_tree_async(group: AsyncGroup, max_depth: int | None = None) -> TreeRepr:
-    tree = rich.tree.Tree(label=f"[b]{group.name}[/b]")
+    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)])
 
@@ -36,9 +43,9 @@ async def group_tree_async(group: AsyncGroup, max_depth: int | None = None) -> T
         # /'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]"
+            label = f"[bold]{name}[/bold]"
         else:
-            label = f"[b]{name}[/b] {node.shape} {node.dtype}"
+            label = f"[bold]{name}[/bold] {node.shape} {node.dtype}"
         nodes[key] = parent.add(label)
 
     return TreeRepr(tree)

tests/test_tree.py

@@ -5,6 +5,8 @@
 
 import zarr
 
+pytest.importorskip("rich")
+
 
 @pytest.mark.parametrize("root_name", [None, "root"])
 def test_tree(root_name: Any) -> None:
@@ -23,26 +25,29 @@ def test_tree(root_name: Any) -> None:
     result = repr(g.tree())
     root = root_name or ""
 
+    BOPEN = "\x1b[1m"
+    BCLOSE = "\x1b[0m"
+
     expected = textwrap.dedent(f"""\
-        /{root}
-        ├── A
-        │   ├── x (2,) float64
-        │   └── y (0,) int8
-        └── B
-            ├── C
-            │   ├── C
-            │   │   └── x (0,) float64
-            │   └── x (0,) float64
-            └── x (0,) float64
+        {BOPEN}/{root}{BCLOSE}
+        ├── {BOPEN}A{BCLOSE}
+        │   ├── {BOPEN}x{BCLOSE} (2,) float64
+        │   └── {BOPEN}y{BCLOSE} (0,) int8
+        └── {BOPEN}B{BCLOSE}
+            ├── {BOPEN}C{BCLOSE}
+            │   ├── {BOPEN}C{BCLOSE}
+            │   │   └── {BOPEN}x{BCLOSE} (0,) float64
+            │   └── {BOPEN}x{BCLOSE} (0,) float64
+            └── {BOPEN}x{BCLOSE} (0,) float64
         """)
 
     assert result == expected
 
     result = repr(g.tree(level=0))
     expected = textwrap.dedent(f"""\
-        /{root}
-        ├── A
-        └── B
+        {BOPEN}/{root}{BCLOSE}
+        ├── {BOPEN}A{BCLOSE}
+        └── {BOPEN}B{BCLOSE}
         """)
 
     assert result == expected