refactored drawtree.py

Author: ramalho

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

@@ -8,15 +8,15 @@
 
 
 def render_lines(tree_iter):
-    name, _, _ = next(tree_iter)
-    yield name
+    cls, _, _ = next(tree_iter)
+    yield cls.__name__
     prefix = ''
 
-    for name, level, last in tree_iter:
+    for cls, 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
+        yield prefix + cls.__name__
 
 
 def draw(cls):

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

@@ -4,7 +4,8 @@
 def test_1_level():
     class One: pass
     expected = [('One', 0, True)]
-    result = list(tree(One))
+    result = [(cls.__name__, level, last)
+              for cls, level, last in tree(One)]
     assert expected == result
 
 
@@ -17,7 +18,8 @@ class Leaf2(Branch): pass
             ('Leaf1', 1, False),
             ('Leaf2', 1, True),
     ]
-    result = list(tree(Branch))
+    result = [(cls.__name__, level, last)
+              for cls, level, last in tree(Branch)]
     assert expected == result
 
 
@@ -30,7 +32,8 @@ class Z(Y): pass
             ('Y', 1, True),
                 ('Z', 2, True),
     ]
-    result = list(tree(X))
+    result = [(cls.__name__, level, last)
+              for cls, level, last in tree(X)]
     assert expected == result
 
 
@@ -46,7 +49,8 @@ class Level3(Level2): pass
                     ('Level3', 3, True),
     ]
 
-    result = list(tree(Level0))
+    result = [(cls.__name__, level, last)
+              for cls, level, last in tree(Level0)]
     assert expected == result
 
 
@@ -68,7 +72,8 @@ class D2(C1): pass
                 ('C2', 2, True),
     ]
 
-    result = list(tree(A))
+    result = [(cls.__name__, level, last)
+              for cls, level, last in tree(A)]
     assert expected == result
 
 
@@ -83,7 +88,8 @@ class Root: pass
         expected.append((name, level, True))
         parent = cls
 
-    result = list(tree(Root))
+    result = [(cls.__name__, level, last)
+              for cls, level, last in tree(Root)]
     assert len(result) == level_count
     assert result[0] == ('Root', 0, True)
     assert result[-1] == ('Sub99', 99, True)

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

@@ -1,5 +1,5 @@
-def tree(cls, level=0, last_in_level=True):
-    yield cls.__name__, level, last_in_level
+def tree(cls, level=0, last_sibling=True):
+    yield cls, level, last_sibling
     subclasses = cls.__subclasses__()
     if subclasses:
         last = subclasses[-1]
@@ -8,9 +8,9 @@ def tree(cls, level=0, last_in_level=True):
 
 
 def display(cls):
-    for cls_name, level, _ in tree(cls):
+    for cls, level, _ in tree(cls):
         indent = ' ' * 4 * level
-        print(f'{indent}{cls_name}')
+        print(f'{indent}{cls.__name__}')
 
 
 if __name__ == '__main__':