From 07b175633eba30dbfcd6eb0cf514ef1b1da9cf64 Mon Sep 17 00:00:00 2001 From: Stephan Hoyer Date: Wed, 26 Jun 2024 11:05:23 -0700 Subject: [PATCH] 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 --- xarray/core/datatree.py | 11 ++++--- xarray/core/datatree_render.py | 11 ++++--- xarray/core/formatting.py | 15 +++------ xarray/core/iterators.py | 19 +++++------ xarray/tests/test_datatree.py | 57 +++++++++++++++++++++++++++++++++ xarray/tests/test_formatting.py | 18 ++++++----- 6 files changed, 94 insertions(+), 37 deletions(-) diff --git a/xarray/core/datatree.py b/xarray/core/datatree.py index 4e4d30885a3..c923ca2eb87 100644 --- a/xarray/core/datatree.py +++ b/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') + + Group: / + ├── Group: /a + │ └── Group: /a/B + └── Group: /b + └── Group: /b/B """ matching_nodes = { node.path: node.ds diff --git a/xarray/core/datatree_render.py b/xarray/core/datatree_render.py index d069071495e..f10f2540952 100644 --- a/xarray/core/datatree_render.py +++ b/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') + + Group: / + ├── Group: /sub0 + │ ├── Group: /sub0/sub0B + │ └── Group: /sub0/sub0A + └── Group: /sub1 """ super().__init__("\u2502 ", "\u251c\u2500\u2500 ", "\u2514\u2500\u2500 ") diff --git a/xarray/core/formatting.py b/xarray/core/formatting.py index ad65a44d7d5..c15df34b5b1 100644 --- a/xarray/core/formatting.py +++ b/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"" + + 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) diff --git a/xarray/core/iterators.py b/xarray/core/iterators.py index dd5fa7ee97a..ae748b0066c 100644 --- a/xarray/core/iterators.py +++ b/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') + + 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)] diff --git a/xarray/tests/test_datatree.py b/xarray/tests/test_datatree.py index 58fec20d4c6..b0dc2accd3e 100644 --- a/xarray/tests/test_datatree.py +++ b/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( + """ + + 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( + """ + + 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}) diff --git a/xarray/tests/test_formatting.py b/xarray/tests/test_formatting.py index b9d5f401a4a..d7a46eeaefc 100644 --- a/xarray/tests/test_formatting.py +++ b/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 == "\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) + + 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)", + "", + "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]})