Skip to content

Commit

Permalink
DAS-2067 - Implement PR feedback, fix RenderDataTree.__str__.
Browse files Browse the repository at this point in the history
  • Loading branch information
owenlittlejohns committed Apr 24, 2024
1 parent e4b7a42 commit 63e2836
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 25 deletions.
17 changes: 9 additions & 8 deletions xarray/core/datatree_render.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,15 +50,14 @@ def __init__(self):
Continued style, without gaps.
>>> from xarray.core.datatree import DataTree
>>> from xarray.core.render import RenderDataTree
>>> from xarray.core.datatree_render import RenderDataTree
>>> root = DataTree(name="root")
>>> s0 = DataTree(name="sub0", parent=root)
>>> s0b = DataTree(name="sub0B", parent=s0)
>>> s0a = DataTree(name="sub0A", parent=s0)
>>> s1 = DataTree(name="sub1", parent=root)
>>> print(RenderDataTree(root))
DataTree('/root', parent=None)
DataTree('root', parent=None)
├── DataTree('sub0')
│ ├── DataTree('sub0B')
│ └── DataTree('sub0A')
Expand Down Expand Up @@ -135,15 +134,18 @@ def __init__(
│ g
└── sub1
h
:any:`by_attr` simplifies attribute rendering and supports multiline:
>>> print(RenderTree(root).by_attr())
>>> print(RenderDataTree(root).by_attr())
root
├── sub0
│ ├── sub0B
│ └── sub0A
└── sub1
`maxlevel` limits the depth of the tree:
>>> print(RenderTree(root, maxlevel=1).by_attr("name"))
# `maxlevel` limits the depth of the tree:
>>> print(RenderDataTree(root, maxlevel=2).by_attr("name"))
root
├── sub0
└── sub1
Expand Down Expand Up @@ -184,8 +186,7 @@ def __item(
return Row(pre, fill, node)

def __str__(self) -> str:
lines = [f"{pre}{node!r}" for pre, _, node in self]
return "\n".join(lines)
return str(self.node)

def __repr__(self) -> str:
classname = self.__class__.__name__
Expand Down
37 changes: 20 additions & 17 deletions xarray/testing/assertions.py
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@ def assert_isomorphic(a: DataTree, b: DataTree, from_root: bool = False):
"""
Two DataTrees are considered isomorphic if every node has the same number of children.
Nothing about the data in each node is checked.
Nothing about the data or attrs in each node is checked.
Isomorphism is a necessary condition for two trees to be used in a nodewise binary operation,
such as tree1 + tree2.
Expand Down Expand Up @@ -96,28 +96,29 @@ def assert_isomorphic(a: DataTree, b: DataTree, from_root: bool = False):


@ensure_warnings
def assert_equal(a, b, **kwargs):
def assert_equal(a, b, from_root=True):
"""Like :py:func:`numpy.testing.assert_array_equal`, but for xarray
objects.
Raises an AssertionError if two objects are not equal. This will match
data values, dimensions and coordinates, but not names or attributes
(except for Dataset objects for which the variable names must match).
Arrays with NaN in the same location are considered equal. DataTrees are
equal if they have isomorphic node structures, with matching node names,
and if they have matching variables and coordinates, all of which are
equal.
Arrays with NaN in the same location are considered equal.
For DataTree objects, assert_equal is mapped over all Datasets on each node,
with the DataTrees being equal if both are isomorphic and the corresponding
Datasets at each node are themselves equal.
Parameters
----------
a : xarray.Dataset, xarray.DataArray, xarray.Variable, xarray.Coordinates
or xarray.core.datatree.DataTree. The first object to compare.
b : xarray.Dataset, xarray.DataArray, xarray.Variable, xarray.Coordinates
or xarray.core.datatree.DataTree. The second object to compare.
**kwargs : dict
Additional keyword arguments passed on to DataTree.equals. Specifically,
from_root, which defaults True, and will check the whole tree above
the given node.
from_root : bool, optional, default is True
Only used when comparing DataTree objects. Indicates whether or not to
first traverse to the root of the trees before checking for isomorphism.
If a & b have no parents then this has no effect.
See Also
--------
Expand All @@ -135,7 +136,6 @@ def assert_equal(a, b, **kwargs):
elif isinstance(a, Coordinates):
assert a.equals(b), formatting.diff_coords_repr(a, b, "equals")
elif isinstance(a, DataTree):
from_root = kwargs.get("from_root", True)
if from_root:
a = a.root
b = b.root
Expand All @@ -146,22 +146,26 @@ def assert_equal(a, b, **kwargs):


@ensure_warnings
def assert_identical(a, b, **kwargs):
def assert_identical(a, b, from_root=True):
"""Like :py:func:`xarray.testing.assert_equal`, but also matches the
objects' names and attributes.
Raises an AssertionError if two objects are not identical.
For DataTree objects, assert_identical is mapped over all Datasets on each
node, with the DataTrees being identical if both are isomorphic and the
corresponding Datasets at each node are themselves identical.
Parameters
----------
a : xarray.Dataset, xarray.DataArray, xarray.Variable or xarray.Coordinates
The first object to compare.
b : xarray.Dataset, xarray.DataArray, xarray.Variable or xarray.Coordinates
The second object to compare.
**kwargs : dict
Additional keyword arguments passed on to DataTree.identical. Specifically,
from_root, which defaults True, and will check the whole tree above
the given node.
from_root : bool, optional, default is True
Only used when comparing DataTree objects. Indicates whether or not to
first traverse to the root of the trees before checking for isomorphism.
If a & b have no parents then this has no effect.
See Also
--------
Expand All @@ -181,7 +185,6 @@ def assert_identical(a, b, **kwargs):
elif isinstance(a, Coordinates):
assert a.identical(b), formatting.diff_coords_repr(a, b, "identical")
elif isinstance(a, DataTree):
from_root = kwargs.get("from_root", True)
if from_root:
a = a.root
b = b.root
Expand Down

0 comments on commit 63e2836

Please sign in to comment.