Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Slightly improve DataTree repr #9064

Merged
merged 8 commits into from
Jun 26, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 6 additions & 5 deletions xarray/core/datatree.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
11 changes: 6 additions & 5 deletions xarray/core/datatree_render.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 ")

Expand Down
15 changes: 5 additions & 10 deletions xarray/core/formatting.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand All @@ -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)


Expand Down
19 changes: 10 additions & 9 deletions xarray/core/iterators.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand Down
57 changes: 57 additions & 0 deletions xarray/tests/test_datatree.py
Original file line number Diff line number Diff line change
Expand Up @@ -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})
Expand Down
18 changes: 10 additions & 8 deletions xarray/tests/test_formatting.py
Original file line number Diff line number Diff line change
Expand Up @@ -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*
Expand All @@ -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",
Expand All @@ -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]})
Expand Down
Loading