refactored drawtree.py

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,35 +0,0 @@
-from tree import tree
-
-SPACES = ' ' * 4
-HLINE  = '\u2500'            # ─ BOX DRAWINGS LIGHT HORIZONTAL
-HLINE2 = HLINE * 2
-ELBOW  = f'\u2514{HLINE2} '  # └ BOX DRAWINGS LIGHT UP AND RIGHT
-TEE    = f'\u251C{HLINE2} '  # ├ BOX DRAWINGS LIGHT VERTICAL AND RIGHT
-PIPE   = '\u2502   '         # │ BOX DRAWINGS LIGHT VERTICAL
-
-
-def render_lines(tree_iter):
-    name, _, _ = next(tree_iter)
-    yield name
-    prefix = ''
-
-    for name, level, last in tree_iter:
-        if last:
-            connector = ELBOW
-        else:
-            connector = TEE
-
-        prefix = prefix[:4 * (level-1)]
-        prefix = prefix.replace(TEE, PIPE).replace(ELBOW, SPACES)
-        prefix += connector
-
-        yield prefix + name
-
-
-def display(cls):
-    for line in render_lines(tree(cls)):
-        print(line)
-
-
-if __name__ == '__main__':
-    display(BaseException)

17-it-generator/tree/extra/test_drawtree.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 @@ def test_2_levels_2_leaves():
     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 @@ def test_3_levels_2_leaves():
     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