diff --git a/datatree/tests/test_treenode.py b/datatree/tests/test_treenode.py index 1805f038..aaa0362a 100644 --- a/datatree/tests/test_treenode.py +++ b/datatree/tests/test_treenode.py @@ -225,57 +225,106 @@ def test_del_child(self): def create_test_tree(): - f = NamedNode() + a = NamedNode(name="a") b = NamedNode() - a = NamedNode() - d = NamedNode() c = NamedNode() + d = NamedNode() e = NamedNode() + f = NamedNode() g = NamedNode() - i = NamedNode() h = NamedNode() + i = NamedNode() - f.children = {"b": b, "g": g} - b.children = {"a": a, "d": d} - d.children = {"c": c, "e": e} - g.children = {"i": i} - i.children = {"h": h} + a.children = {"b": b, "c": c} + b.children = {"d": d, "e": e} + e.children = {"f": f, "g": g} + c.children = {"h": h} + h.children = {"i": i} - return f + return a, f class TestIterators: def test_preorderiter(self): - tree = create_test_tree() - result = [node.name for node in PreOrderIter(tree)] + root, _ = create_test_tree() + result = [node.name for node in PreOrderIter(root)] expected = [ - None, # root Node is unnamed - "b", "a", + "b", "d", - "c", "e", + "f", "g", - "i", + "c", "h", + "i", ] assert result == expected def test_levelorderiter(self): - tree = create_test_tree() - result = [node.name for node in LevelOrderIter(tree)] + root, _ = create_test_tree() + result = [node.name for node in LevelOrderIter(root)] expected = [ - None, # root Node is unnamed + "a", # root Node is unnamed "b", + "c", + "d", + "e", + "h", + "f", "g", + "i", + ] + assert result == expected + + +class TestAncestry: + def test_lineage(self): + _, leaf = create_test_tree() + lineage = leaf.lineage + expected = ["f", "e", "b", "a"] + for node, expected_name in zip(lineage, expected): + assert node.name == expected_name + + def test_ancestors(self): + _, leaf = create_test_tree() + ancestors = leaf.ancestors + expected = ["a", "b", "e", "f"] + for node, expected_name in zip(ancestors, expected): + assert node.name == expected_name + + def test_subtree(self): + root, _ = create_test_tree() + subtree = root.subtree + expected = [ "a", + "b", "d", - "i", + "e", + "f", + "g", "c", + "h", + "i", + ] + for node, expected_name in zip(subtree, expected): + assert node.name == expected_name + + def test_descendants(self): + root, _ = create_test_tree() + descendants = root.descendants + expected = [ + "b", + "d", "e", + "f", + "g", + "c", "h", + "i", ] - assert result == expected + for node, expected_name in zip(descendants, expected): + assert node.name == expected_name class TestRenderTree: diff --git a/datatree/treenode.py b/datatree/treenode.py index 3ff2eb98..7b499710 100644 --- a/datatree/treenode.py +++ b/datatree/treenode.py @@ -238,7 +238,6 @@ def _post_attach_children(self: Tree, children: Mapping[str, Tree]) -> None: def iter_lineage(self: Tree) -> Iterator[Tree]: """Iterate up the tree, starting from the current node.""" - # TODO should this instead return an OrderedDict, so as to include node names? node: Tree | None = self while node is not None: yield node @@ -298,11 +297,30 @@ def subtree(self: Tree) -> Iterator[Tree]: An iterator over all nodes in this tree, including both self and all descendants. Iterates depth-first. + + See Also + -------- + DataTree.descendants """ from . import iterators return iterators.PreOrderIter(self) + @property + def descendants(self: Tree) -> Tuple[Tree]: + """ + Child nodes and all their child nodes. + + Returned in depth-first order. + + See Also + -------- + DataTree.subtree + """ + all_nodes = tuple(self.subtree) + this_node, *descendants = all_nodes + return tuple(descendants) # type: ignore[return-value] + def _pre_detach(self: Tree, parent: Tree) -> None: """Method call before detaching from `parent`.""" pass diff --git a/docs/source/api.rst b/docs/source/api.rst index 75c70584..751e5643 100644 --- a/docs/source/api.rst +++ b/docs/source/api.rst @@ -31,6 +31,7 @@ Attributes relating to the recursive tree-like structure of a ``DataTree``. DataTree.is_root DataTree.is_leaf DataTree.subtree + DataTree.descendants DataTree.siblings DataTree.lineage DataTree.ancestors diff --git a/docs/source/whats-new.rst b/docs/source/whats-new.rst index 18ac3024..5e86ba1c 100644 --- a/docs/source/whats-new.rst +++ b/docs/source/whats-new.rst @@ -28,6 +28,9 @@ New Features - Allow method chaining with a new :py:meth:`DataTree.pipe` method (:issue:`151`, :pull:`156`). By `Justus Magin `_. - New, more specific exception types for tree-related errors (:pull:`169`). + By `Tom Nicholas `_. +- Added a new :py:meth:`DataTree.descendants` property (:pull:`170`). + By `Tom Nicholas `_. Breaking changes ~~~~~~~~~~~~~~~~