From 40b41490648eb52dfbe7a2b8f1e278a9c4e0c4ce Mon Sep 17 00:00:00 2001 From: Giuseppe Corbelli Date: Thu, 12 Sep 2019 10:23:51 +0200 Subject: [PATCH 1/2] Added a Tree.children_itr method to be a companion to Tree.children in case an iterator is preferrable over a list creation for efficiency reasons. --- treelib/tree.py | 7 +++++++ 1 file changed, 7 insertions(+) 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 From d3ff7e1e561f00b8d9fb578ad66dd69db1ce3d46 Mon Sep 17 00:00:00 2001 From: Giuseppe Corbelli Date: Thu, 12 Sep 2019 10:27:11 +0200 Subject: [PATCH 2/2] Added a simple test for Tree.children_itr(). It must return the same nodes as its companion method Tree.children(). --- tests/test_tree.py | 2 ++ 1 file changed, 2 insertions(+) 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: