Slightly improve DataTree repr (#9064)

* Improve DataTree repr

* Adjust DataTree repr to include full path

* More tweaks

* Use "Group:" in repr instead of "DataTree:"

* Fix errors in new repr tests

* Fix repr on windows

Author: shoyer

xarray/core/datatree.py

@@ -1314,11 +1314,12 @@ def match(self, pattern: str) -> DataTree:
         ...     }
         ... )
         >>> dt.match("*/B")
-        DataTree('None', parent=None)
-        ├── DataTree('a')
-        │   └── DataTree('B')
-        └── DataTree('b')
-            └── DataTree('B')
+        <xarray.DataTree>
+        Group: /
+        ├── Group: /a
+        │   └── Group: /a/B
+        └── Group: /b
+            └── Group: /b/B
         """
         matching_nodes = {
             node.path: node.ds

xarray/core/datatree_render.py

@@ -57,11 +57,12 @@ def __init__(self):
         >>> s0a = DataTree(name="sub0A", parent=s0)
         >>> s1 = DataTree(name="sub1", parent=root)
         >>> print(RenderDataTree(root))
-        DataTree('root', parent=None)
-        ├── DataTree('sub0')
-        │   ├── DataTree('sub0B')
-        │   └── DataTree('sub0A')
-        └── DataTree('sub1')
+        <xarray.DataTree 'root'>
+        Group: /
+        ├── Group: /sub0
+        │   ├── Group: /sub0/sub0B
+        │   └── Group: /sub0/sub0A
+        └── Group: /sub1
         """
         super().__init__("\u2502   ", "\u251c\u2500\u2500 ", "\u2514\u2500\u2500 ")
 

xarray/core/formatting.py

@@ -1023,20 +1023,21 @@ def diff_datatree_repr(a: DataTree, b: DataTree, compat):
 
 def _single_node_repr(node: DataTree) -> str:
     """Information about this node, not including its relationships to other nodes."""
-    node_info = f"DataTree('{node.name}')"
-
     if node.has_data or node.has_attrs:
         ds_info = "\n" + repr(node.ds)
     else:
         ds_info = ""
-    return node_info + ds_info
+    return f"Group: {node.path}{ds_info}"
 
 
 def datatree_repr(dt: DataTree):
     """A printable representation of the structure of this entire tree."""
     renderer = RenderDataTree(dt)
 
-    lines = []
+    name_info = "" if dt.name is None else f" {dt.name!r}"
+    header = f"<xarray.DataTree{name_info}>"
+
+    lines = [header]
     for pre, fill, node in renderer:
         node_repr = _single_node_repr(node)
 
@@ -1051,12 +1052,6 @@ def datatree_repr(dt: DataTree):
                 else:
                     lines.append(f"{fill}{' ' * len(renderer.style.vertical)}{line}")
 
-    # Tack on info about whether or not root node has a parent at the start
-    first_line = lines[0]
-    parent = f'"{dt.parent.name}"' if dt.parent is not None else "None"
-    first_line_with_parent = first_line[:-1] + f", parent={parent})"
-    lines[0] = first_line_with_parent
-
     return "\n".join(lines)
 
 

xarray/core/iterators.py

@@ -39,15 +39,16 @@ class LevelOrderIter(Iterator):
     >>> i = DataTree(name="i", parent=g)
     >>> h = DataTree(name="h", parent=i)
     >>> print(f)
-    DataTree('f', parent=None)
-    ├── DataTree('b')
-    │   ├── DataTree('a')
-    │   └── DataTree('d')
-    │       ├── DataTree('c')
-    │       └── DataTree('e')
-    └── DataTree('g')
-        └── DataTree('i')
-            └── DataTree('h')
+    <xarray.DataTree 'f'>
+    Group: /
+    ├── Group: /b
+    │   ├── Group: /b/a
+    │   └── Group: /b/d
+    │       ├── Group: /b/d/c
+    │       └── Group: /b/d/e
+    └── Group: /g
+        └── Group: /g/i
+            └── Group: /g/i/h
     >>> [node.name for node in LevelOrderIter(f)]
     ['f', 'b', 'g', 'a', 'd', 'i', 'c', 'e', 'h']
     >>> [node.name for node in LevelOrderIter(f, maxlevel=3)]

xarray/tests/test_datatree.py

@@ -623,6 +623,63 @@ def test_operation_with_attrs_but_no_data(self):
         dt.sel(dim_0=0)
 
 
+class TestRepr:
+    def test_repr(self):
+        dt: DataTree = DataTree.from_dict(
+            {
+                "/": xr.Dataset(
+                    {"e": (("x",), [1.0, 2.0])},
+                    coords={"x": [2.0, 3.0]},
+                ),
+                "/b": xr.Dataset({"f": (("y",), [3.0])}),
+                "/b/c": xr.Dataset(),
+                "/b/d": xr.Dataset({"g": 4.0}),
+            }
+        )
+
+        result = repr(dt)
+        expected = dedent(
+            """
+            <xarray.DataTree>
+            Group: /
+            │   Dimensions:  (x: 2)
+            │   Coordinates:
+            │     * x        (x) float64 16B 2.0 3.0
+            │   Data variables:
+            │       e        (x) float64 16B 1.0 2.0
+            └── Group: /b
+                │   Dimensions:  (y: 1)
+                │   Dimensions without coordinates: y
+                │   Data variables:
+                │       f        (y) float64 8B 3.0
+                ├── Group: /b/c
+                └── Group: /b/d
+                        Dimensions:  ()
+                        Data variables:
+                            g        float64 8B 4.0
+            """
+        ).strip()
+        assert result == expected
+
+        result = repr(dt.b)
+        expected = dedent(
+            """
+            <xarray.DataTree 'b'>
+            Group: /b
+            │   Dimensions:  (y: 1)
+            │   Dimensions without coordinates: y
+            │   Data variables:
+            │       f        (y) float64 8B 3.0
+            ├── Group: /b/c
+            └── Group: /b/d
+                    Dimensions:  ()
+                    Data variables:
+                        g        float64 8B 4.0
+            """
+        ).strip()
+        assert result == expected
+
+
 class TestRestructuring:
     def test_drop_nodes(self):
         sue = DataTree.from_dict({"Mary": None, "Kate": None, "Ashley": None})

xarray/tests/test_formatting.py

@@ -555,16 +555,17 @@ def test_array_scalar_format(self) -> None:
 
     def test_datatree_print_empty_node(self):
         dt: DataTree = DataTree(name="root")
-        printout = dt.__str__()
-        assert printout == "DataTree('root', parent=None)"
+        printout = str(dt)
+        assert printout == "<xarray.DataTree 'root'>\nGroup: /"
 
     def test_datatree_print_empty_node_with_attrs(self):
         dat = xr.Dataset(attrs={"note": "has attrs"})
         dt: DataTree = DataTree(name="root", data=dat)
-        printout = dt.__str__()
+        printout = str(dt)
         assert printout == dedent(
             """\
-            DataTree('root', parent=None)
+            <xarray.DataTree 'root'>
+            Group: /
                 Dimensions:  ()
                 Data variables:
                     *empty*
@@ -575,9 +576,10 @@ def test_datatree_print_empty_node_with_attrs(self):
     def test_datatree_print_node_with_data(self):
         dat = xr.Dataset({"a": [0, 2]})
         dt: DataTree = DataTree(name="root", data=dat)
-        printout = dt.__str__()
+        printout = str(dt)
         expected = [
-            "DataTree('root', parent=None)",
+            "<xarray.DataTree 'root'>",
+            "Group: /",
             "Dimensions",
             "Coordinates",
             "a",
@@ -591,8 +593,8 @@ def test_datatree_printout_nested_node(self):
         dat = xr.Dataset({"a": [0, 2]})
         root: DataTree = DataTree(name="root")
         DataTree(name="results", data=dat, parent=root)
-        printout = root.__str__()
-        assert printout.splitlines()[2].startswith("    ")
+        printout = str(root)
+        assert printout.splitlines()[3].startswith("    ")
 
     def test_datatree_repr_of_node_with_data(self):
         dat = xr.Dataset({"a": [0, 2]})