ch17: update from book draft

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

@@ -0,0 +1,38 @@
+from tree import tree
+
+
+def test_1_level():
+    class One: pass
+    expected = [('One', 0)]
+    result = list(tree(One))
+    assert expected == result 
+
+
+def test_2_levels_4_leaves():
+    class Branch: pass
+    class Leaf1(Branch): pass
+    class Leaf2(Branch): pass
+    class Leaf3(Branch): pass
+    class Leaf4(Branch): pass
+    expected = [
+        ('Branch', 0),
+            ('Leaf1', 1),
+            ('Leaf2', 1),
+            ('Leaf3', 1),
+            ('Leaf4', 1),
+    ]
+    result = list(tree(Branch))
+    assert expected == result 
+
+
+def test_3_levels_1_leaf():
+    class X: pass
+    class Y(X): pass
+    class Z(Y): pass
+    expected = [
+        ('X', 0),
+            ('Y', 1),
+                ('Z', 2),
+    ]
+    result = list(tree(X))
+    assert expected == result 

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

@@ -0,0 +1,16 @@
+def tree(cls):
+    yield cls.__name__, 0
+    yield from sub_tree(cls)
+
+
+def sub_tree(cls):
+    for sub_cls in cls.__subclasses__():
+        yield sub_cls.__name__, 1
+        for sub_sub_cls in sub_cls.__subclasses__():
+            yield sub_sub_cls.__name__, 2
+
+
+if __name__ == '__main__':
+    for cls_name, level in tree(BaseException):
+        indent = ' ' * 4 * level
+        print(f'{indent}{cls_name}')