Descendants property

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:

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

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

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 <https://github.com/keewis>`_.
 - New, more specific exception types for tree-related errors (:pull:`169`).
+  By `Tom Nicholas <https://github.com/TomNicholas>`_.
+- Added a new :py:meth:`DataTree.descendants` property (:pull:`170`).
+  By `Tom Nicholas <https://github.com/TomNicholas>`_.
 
 Breaking changes
 ~~~~~~~~~~~~~~~~