Merge pull request #63 from BebeSparkelSparkel/master

Added all_nodes_iter and filter_nodes methods to Tree class

tests/test_plugins.py

@@ -17,14 +17,14 @@ def setUp(self):
         tree.create_node("Diane", "diane", parent="jane")
         tree.create_node("George", "george", parent="bill")
         self.tree = tree
-    
+
     def read_generated_output(self, filename):
         output = codecs.open(filename, 'r', 'utf-8')
         generated = output.read()
         output.close()
-        
+
         return generated
-    
+
     def test_export_to_dot(self):
         export_to_dot(self.tree, 'tree.dot')
         expected = """\
@@ -40,32 +40,32 @@ def test_export_to_dot(self):
 \t"bill" -> "george"
 \t"jane" -> "diane"
 }"""
-        
+
         self.assertTrue(os.path.isfile('tree.dot'), "The file tree.dot could not be found.")
         generated = self.read_generated_output('tree.dot')
-        
-        self.assertEqual(generated, expected, "Generated dot tree is not the expected one")  
+
+        self.assertEqual(generated, expected, "Generated dot tree is not the expected one")
         os.remove('tree.dot')
-        
+
     def test_export_to_dot_empty_tree(self):
         empty_tree = Tree()
         export_to_dot(empty_tree, 'tree.dot')
-        
+
         expected = """\
 digraph tree {
 
 }"""
         self.assertTrue(os.path.isfile('tree.dot'), "The file tree.dot could not be found.")
         generated = self.read_generated_output('tree.dot')
-        
+
         self.assertEqual(expected, generated, 'The generated output for an empty tree is not empty')
         os.remove('tree.dot')
-    
+
     def test_unicode_filename(self):
         tree = Tree()
         tree.create_node('Node 1', 'node_1')
         export_to_dot(tree, 'ŕʩϢ.dot')
-        
+
         expected = """\
 digraph tree {
 \t"node_1" [label="Node 1", shape=circle]

tests/test_treelib.py

@@ -318,6 +318,37 @@ def tearDown(self):
         self.tree = None
         self.copytree = None
 
+    def test_all_nodes_itr(self):
+        """
+        tests: Tree.all_nodes_iter
+        Added by: William Rusnack
+        """
+        new_tree = Tree()
+        self.assertEqual(len(new_tree.all_nodes_itr()), 0)
+        nodes = []
+        nodes.append(new_tree.create_node('root_node'))
+        nodes.append(new_tree.create_node('second', parent=new_tree.root))
+        for nd in new_tree.all_nodes_itr():
+            self.assertTrue(nd in nodes)
+
+    def test_filter_nodes(self):
+        """
+        tests: Tree.filter_nodes
+        Added by: William Rusnack
+        """
+        new_tree = Tree()
+
+        self.assertEqual(tuple(new_tree.filter_nodes(lambda n: True)), ())
+
+        nodes = []
+        nodes.append(new_tree.create_node('root_node'))
+        nodes.append(new_tree.create_node('second', parent=new_tree.root))
+
+        self.assertEqual(tuple(new_tree.filter_nodes(lambda n: False)), ())
+        self.assertEqual(tuple(new_tree.filter_nodes(lambda n: n.is_root())), (nodes[0],))
+        self.assertEqual(tuple(new_tree.filter_nodes(lambda n: not n.is_root())), (nodes[1],))
+        self.assertTrue(set(new_tree.filter_nodes(lambda n: True)), set(nodes))
+
 def suite():
     suites = [NodeCase, TreeCase]
     suite = unittest.TestSuite()

treelib/__init__.py

@@ -1,4 +1,4 @@
-__version__ = '1.3.6'
+__version__ = '1.3.7'
 
 from .tree import Tree
 from .node import Node

treelib/tree.py

@@ -275,6 +275,13 @@ def all_nodes(self):
         """Return all nodes in a list"""
         return list(self._nodes.values())
 
+    def all_nodes_itr(self):
+        """
+        Returns all nodes in a iterator
+        Added by William Rusnack
+        """
+        return self._nodes.values()
+
     def children(self, nid):
         """
         Return the children (Node) list of nid.
@@ -610,6 +617,15 @@ def rsearch(self, nid, filter=None):
             # subtree() hasn't update the bpointer
             current = self[current].bpointer if self.root != current else None
 
+    def filter_nodes(self, function):
+        """
+        Filters all nodes by function
+        function is passed one node as an argument and that node is included if function returns true
+        returns a filter iterator of the node in python 3 or a list of the nodes in python 2
+        Added William Rusnack
+        """
+        return filter(function, self.all_nodes_itr())
+
     def save2file(self, filename, nid=None, level=ROOT, idhidden=True,
                   filter=None, key=None, reverse=False, line_type='ascii-ex', data_property=None):
         """Update 20/05/13: Save tree into file for offline analysis"""