diff --git a/tests/test_tree.py b/tests/test_tree.py index 3ca55c0..74ac8e3 100644 --- a/tests/test_tree.py +++ b/tests/test_tree.py @@ -97,6 +97,8 @@ def test_children(self): children = self.tree.children(nid) for child in children: self.assertEqual(child in self.tree.all_nodes(), True) + children_from_itr = list(self.tree.children_itr(nid)) + self.assertEqual(children, children_from_itr) try: self.tree.is_branch("alien") except NodeIDAbsentError: diff --git a/treelib/tree.py b/treelib/tree.py index 62d3549..fd56041 100644 --- a/treelib/tree.py +++ b/treelib/tree.py @@ -296,6 +296,13 @@ def children(self, nid): """ return [self[i] for i in self.is_branch(nid)] + def children_itr(self, nid): + """ + Return the children (Node) of nid as an iterator. + Empty iterator is returned if nid does not exist + """ + return (self[i] for i in self.is_branch(nid)) + def contains(self, nid): """Check if the tree contains node of given id""" return True if nid in self._nodes else False