flatten testswq

2 files changed, 23 insertions, 1 deletions

diff --git a/test/unittest_tree.py b/test/unittest_tree.py
index aecb284..598a4a2 100644
--- a/test/unittest_tree.py
+++ b/test/unittest_tree.py
@@ -24,6 +24,21 @@ class Node_ClassTest(TestCase):
     def setUp(self):
         """ called before each test from this class """        
         self.o = make_tree(tree)
+
+
+    def test_flatten(self):
+        result = [r.id for r in self.o.flatten()]
+        expected = ['root', 'child_1_1', 'child_2_1', 'child_2_2', 'child_3_1', 'child_1_2', 'child_2_3']
+        self.assertListEqual(result, expected)
+
+    def test_flatten_with_outlist(self):
+        resultnodes = []
+        self.o.flatten(resultnodes)
+        result = [r.id for r in resultnodes]
+        expected = ['root', 'child_1_1', 'child_2_1', 'child_2_2', 'child_3_1', 'child_1_2', 'child_2_3']
+        self.assertListEqual(result, expected)
+        
+        
     def test_known_values_remove(self):
         """ 
         remove a child node
diff --git a/tree.py b/tree.py
index 435b3f0..80f56fc 100644
--- a/tree.py
+++ b/tree.py
@@ -190,7 +190,14 @@ class Node :
         """
         return a list with all the nodes descendant from this node
         """
-        return flatten(self)
+        if _list is None:
+            _list = []
+        _list.append(self)
+        for c in self.children:
+            c.flatten(_list)
+        return _list
+
+##        return flatten(self)
 
     def lineage(self):
         """