bugfix/enchancement for prettyTree

yoyzhou · Sep 8, 2012 · ce0e6a1 · ce0e6a1
1 parent 0fce493
commit ce0e6a1
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 16 deletions.
diff --git a/src/pyTree/Tree.py b/src/pyTree/Tree.py
@@ -220,7 +220,7 @@ def prettyTree(self):
             if isinstance(head, int):
                 level = head
             else:
-                self.__printLabel__(str(head.data), level)
+                self.__printLabel__(head, NodesS, level)
                 children = head.getChildren()
                 children.reverse()
 
@@ -264,25 +264,39 @@ def nestedTree(self):
 
         print(NestedT)
 
-    def __printLabel__(self, label, level):
+    def __printLabel__(self, head, NodesS, level):
         """
            Print each node
         """
+        leading = '' 
+        lasting = '|___ '
+        label = str(head)
+
         if level == 0:
-            print(label)
+            print(str(head))
         else:
-            if level == 1:
-                leadingSpaces = ''
-                addl = ''
-            else:
-                leadingSpaces =   ' ' * 5 * (level - 1)
-                addl = '|'
+            for l in range(0, level - 1):
+                sibling = False
+                parentT = head.__getParent__(level - l)
+                for c in parentT.getChildren():
+                    if c in NodesS:
+                        sibling = True
+                        break
+                if sibling:
+                    leading += '|     '
+                else:
+                    leading += '     '
 
             if label.strip() != '': 
-                print('{0}{1}{2}{3} {4}'.format('|', leadingSpaces, addl, '_' * 3, label))
+                print('{0}{1}{2}'.format( leading, lasting, label))
 
-
-
+
+    def __getParent__(self, up):
+        parent = self;
+        while up:
+            parent = parent.getParent()
+            up -= 1
+        return parent
 
 
 

diff --git a/test/testTree.py b/test/testTree.py
@@ -110,17 +110,34 @@ def test_GetNode(self):
 
     def test_DelChild(self):
         self.child11.delChild(0);
-        self.root.prettyTree()
+        #self.root.prettyTree()
 
 
     def test_DelNode(self):
         self.root.delNode('C03');
-        self.root.prettyTree()
+        #self.root.prettyTree()
 
 
     def test_PrintTree(self):
-        self.root.prettyTree()
-        self.root.nestedTree()
+        root = Tree('root')
+        t1 = Tree('1')
+        t11 = Tree('11')
+        t111 = Tree('111')
+        t1111 = Tree('1111')
+        t11111 = Tree('11111')
+        t111111 = Tree('111111')
+        t1111111 = Tree('1111111')
+        root.addChild(t1)
+        t1.addChild(t11)
+        t11.addChild(t111)
+        t111.addChild(t1111)
+        t1111.addChild(t11111)
+        t1111.addChild(Tree('44444'))
+        t11111.addChild(t111111)
+        t111111.addChild(t1111111)
+        root.prettyTree()       
+        #self.root.prettyTree()
+        # self.root.nestedTree()
 
 
     def tearDown(self):