refactored drawtree.py

Author: ramalho

17-it-generator/tree/extra/drawtree.py

@@ -0,0 +1,28 @@
+from tree import tree
+
+SP = '\N{SPACE}'
+HLIN  = '\N{BOX DRAWINGS LIGHT HORIZONTAL}'                       # ─
+ELBOW = f'\N{BOX DRAWINGS LIGHT UP AND RIGHT}{HLIN*2}{SP}'        # └──
+TEE   = f'\N{BOX DRAWINGS LIGHT VERTICAL AND RIGHT}{HLIN*2}{SP}'  # ├──
+PIPE  = f'\N{BOX DRAWINGS LIGHT VERTICAL}{SP*3}'                  # │
+
+
+def render_lines(tree_iter):
+    name, _, _ = next(tree_iter)
+    yield name
+    prefix = ''
+
+    for name, level, last in tree_iter:
+        prefix = prefix[:4 * (level-1)]
+        prefix = prefix.replace(TEE, PIPE).replace(ELBOW, SP*4)
+        prefix += ELBOW if last else TEE
+        yield prefix + name
+
+
+def draw(cls):
+    for line in render_lines(tree(cls)):
+        print(line)
+
+
+if __name__ == '__main__':
+    draw(BaseException)

17-it-generator/tree/extra/pretty_tree.py

@@ -1,4 +1,4 @@
-from pretty_tree import tree, render_lines
+from drawtree import tree, render_lines
 
 def test_1_level():
     result = list(render_lines(tree(BrokenPipeError)))
@@ -59,7 +59,7 @@ class Leaf2(Branch): pass
     assert expected == result
 
 
-def test_3_levels_2_leaves():
+def test_3_levels_2_leaves_dedent():
     class A: pass
     class B(A): pass
     class C(B): pass
@@ -77,7 +77,7 @@ class E(D): pass
     assert expected == result
 
 
-def test_4_levels_4_leaves():
+def test_4_levels_4_leaves_dedent():
     class A: pass
     class B1(A): pass
     class C1(B1): pass