From c41611668dbc9431370b1b61dfdab80497794432 Mon Sep 17 00:00:00 2001
From: Luciano Ramalho <luciano@ramalho.org>
Date: Mon, 18 Oct 2021 14:57:06 -0300
Subject: [PATCH] refactored drawtree.py

---
 17-it-generator/tree/extra/drawtree.py  |  8 ++++----
 17-it-generator/tree/extra/test_tree.py | 18 ++++++++++++------
 17-it-generator/tree/extra/tree.py      |  8 ++++----
 3 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/17-it-generator/tree/extra/drawtree.py b/17-it-generator/tree/extra/drawtree.py
index 043e53b..522666a 100644
--- a/17-it-generator/tree/extra/drawtree.py
+++ b/17-it-generator/tree/extra/drawtree.py
@@ -8,15 +8,15 @@ PIPE  = f'\N{BOX DRAWINGS LIGHT VERTICAL}{SP*3}'                  # │
 
 
 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):
diff --git a/17-it-generator/tree/extra/test_tree.py b/17-it-generator/tree/extra/test_tree.py
index 4c6404c..b06cf3b 100644
--- a/17-it-generator/tree/extra/test_tree.py
+++ b/17-it-generator/tree/extra/test_tree.py
@@ -4,7 +4,8 @@ from tree import tree
 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 @@ def test_2_levels_2_leaves():
             ('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 @@ def test_3_levels_1_leaf():
             ('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 @@ def test_4_levels_1_leaf():
                     ('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 @@ def test_4_levels_3_leaves():
                 ('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 @@ def test_many_levels_1_leaf():
         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)
diff --git a/17-it-generator/tree/extra/tree.py b/17-it-generator/tree/extra/tree.py
index 7169a28..aeebd37 100644
--- a/17-it-generator/tree/extra/tree.py
+++ b/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__':