minor refactorings

ramalho · ramalho · commit f5e3cb8ad32b · 2021-10-22T14:59:04.000-03:00
diff --git a/17-it-generator/tree/classtree/classtree.py b/17-it-generator/tree/classtree/classtree.py
@@ -0,0 +1,75 @@
+#!/usr/bin/env python3
+
+from importlib import import_module
+import sys
+
+
+SP    = '\N{SPACE}'
+HLIN  = '\N{BOX DRAWINGS LIGHT HORIZONTAL}' * 2 + SP        # ──
+VLIN  = '\N{BOX DRAWINGS LIGHT VERTICAL}' + SP * 3          # │
+TEE   = '\N{BOX DRAWINGS LIGHT VERTICAL AND RIGHT}' + HLIN  # ├──
+ELBOW = '\N{BOX DRAWINGS LIGHT UP AND RIGHT}' + HLIN        # └──
+
+
+def subclasses(cls):
+    try:
+        return cls.__subclasses__()
+    except TypeError:  # handle the `type` type
+        return cls.__subclasses__(cls)
+
+
+def tree(cls, level=0, last_sibling=True):
+    yield cls, level, last_sibling
+    chidren = subclasses(cls)
+    if chidren:
+        last = chidren[-1]
+        for child in chidren:
+            yield from tree(child, level + 1, child is last)
+
+
+def render_lines(tree_generator):
+    cls, _, _ = next(tree_generator)
+    yield cls.__name__
+    prefix = ''
+    for cls, level, last in tree_generator:
+        prefix = prefix[: 4 * (level - 1)]
+        prefix = prefix.replace(TEE, VLIN).replace(ELBOW, SP * 4)
+        prefix += ELBOW if last else TEE
+        yield prefix + cls.__name__
+
+
+def draw(cls):
+    for line in render_lines(tree(cls)):
+        print(line)
+
+
+def parse(name):
+    if '.' in name:
+        return name.rsplit('.', 1)
+    else:
+        return 'builtins', name
+
+
+def main(name):
+    module_name, cls_name = parse(name)
+    try:
+        cls = getattr(import_module(module_name), cls_name)
+    except ModuleNotFoundError:
+        print(f'*** Could not import {module_name!r}.')
+    except AttributeError:
+        print(f'*** {cls_name!r} not found in {module_name!r}.')
+    else:
+        if isinstance(cls, type):
+            draw(cls)
+        else:
+            print(f'*** {cls_name!r} is not a class.')
+
+
+if __name__ == '__main__':
+    if len(sys.argv) == 2:
+        main(sys.argv[1])
+    else:
+        print('Usage:'
+            f'\t{sys.argv[0]} Class          # for builtin classes\n'
+            f'\t{sys.argv[0]} package.Class  # for other classes'
+        )
diff --git a/17-it-generator/tree/classtree/classtree_test.py b/17-it-generator/tree/classtree/classtree_test.py
@@ -0,0 +1,181 @@
+from textwrap import dedent
+from typing import SupportsBytes
+from classtree import tree, render_lines, main, subclasses
+
+from abc import ABCMeta
+
+
+def test_subclasses():
+    result = subclasses(UnicodeError)
+    assert set(result) >= {UnicodeEncodeError, UnicodeDecodeError}
+
+
+def test_subclasses_of_type():
+    """
+    The `type` class is a special case because `type.__subclasses__()`
+    is an unbound method when called on it, so we must call it as
+    `type.__subclasses__(type)` just for `type`.
+
+    This test does not verify the full list of results, but just
+    checks that `abc.ABCMeta` is included, because that's the only
+    subclass of `type` (i.e. metaclass) I we get when I run
+    `$ classtree.py type` at the command line.
+
+    However, the Python console and `pytest` both load other modules,
+    so `subclasses` may find more subclasses of `type`—for example,
+    `enum.EnumMeta`.
+    """
+    result = subclasses(type)
+    assert ABCMeta in result
+
+
+def test_tree_1_level():
+    result = list(tree(TabError))
+    assert result == [(TabError, 0, True)]
+
+
+def test_tree_2_levels():
+    result = list(tree(IndentationError))
+    assert result == [
+        (IndentationError, 0, True),
+        (TabError, 1, True),
+    ]
+
+
+def test_render_lines_1_level():
+    result = list(render_lines(tree(TabError)))
+    assert result == ['TabError']
+
+
+def test_render_lines_2_levels_1_leaf():
+    result = list(render_lines(tree(IndentationError)))
+    expected = [
+        'IndentationError',
+        '└── TabError',
+    ]
+    assert expected == result
+
+
+def test_render_lines_3_levels_1_leaf():
+    class X: pass
+    class Y(X): pass
+    class Z(Y): pass
+    result = list(render_lines(tree(X)))
+    expected = [
+        'X',
+        '└── Y',
+        '    └── Z',
+    ]
+    assert expected == result
+
+
+def test_render_lines_4_levels_1_leaf():
+    class Level0: pass
+    class Level1(Level0): pass
+    class Level2(Level1): pass
+    class Level3(Level2): pass
+
+    result = list(render_lines(tree(Level0)))
+    expected = [
+        'Level0',
+        '└── Level1',
+        '    └── Level2',
+        '        └── Level3',
+    ]
+    assert expected == result
+
+
+def test_render_lines_2_levels_2_leaves():
+    class Branch: pass
+    class Leaf1(Branch): pass
+    class Leaf2(Branch): pass
+    result = list(render_lines(tree(Branch)))
+    expected = [
+        'Branch',
+        '├── Leaf1',
+        '└── Leaf2',
+    ]
+    assert expected == result
+
+
+def test_render_lines_3_levels_2_leaves_dedent():
+    class A: pass
+    class B(A): pass
+    class C(B): pass
+    class D(A): pass
+    class E(D): pass
+
+    result = list(render_lines(tree(A)))
+    expected = [
+        'A',
+        '├── B',
+        '│   └── C',
+        '└── D',
+        '    └── E',
+    ]
+    assert expected == result
+
+
+def test_render_lines_4_levels_4_leaves_dedent():
+    class A: pass
+    class B1(A): pass
+    class C1(B1): pass
+    class D1(C1): pass
+    class D2(C1): pass
+    class C2(B1): pass
+    class B2(A): pass
+    expected = [
+        'A',
+        '├── B1',
+        '│   ├── C1',
+        '│   │   ├── D1',
+        '│   │   └── D2',
+        '│   └── C2',
+        '└── B2',
+    ]
+
+    result = list(render_lines(tree(A)))
+    assert expected == result
+
+
+def test_main_simple(capsys):
+    main('IndentationError')
+    expected = dedent("""
+        IndentationError
+        └── TabError
+    """).lstrip()
+    captured = capsys.readouterr()
+    assert captured.out == expected
+
+
+def test_main_dotted(capsys):
+    main('collections.abc.Sequence')
+    expected = dedent("""
+        Sequence
+        ├── ByteString
+        ├── MutableSequence
+        │   └── UserList
+    """).lstrip()
+    captured = capsys.readouterr()
+    assert captured.out.startswith(expected)
+
+
+def test_main_class_not_found(capsys):
+    main('NoSuchClass')
+    expected = "*** 'NoSuchClass' not found in 'builtins'.\n"
+    captured = capsys.readouterr()
+    assert captured.out == expected
+
+
+def test_main_module_not_found(capsys):
+    main('nosuch.module')
+    expected = "*** Could not import 'nosuch'.\n"
+    captured = capsys.readouterr()
+    assert captured.out == expected
+
+
+def test_main_not_a_class(capsys):
+    main('collections.abc')
+    expected = "*** 'abc' is not a class.\n"
+    captured = capsys.readouterr()
+    assert captured.out == expected
diff --git a/17-it-generator/tree/extra/drawtree.py b/17-it-generator/tree/extra/drawtree.py
@@ -1,10 +1,26 @@
 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}'                  # │
+SP    = '\N{SPACE}'
+HLIN  = '\N{BOX DRAWINGS LIGHT HORIZONTAL}' * 2 + SP        # ──
+VLIN  = '\N{BOX DRAWINGS LIGHT VERTICAL}' + SP * 3          # │
+TEE   = '\N{BOX DRAWINGS LIGHT VERTICAL AND RIGHT}' + HLIN  # ├──
+ELBOW = '\N{BOX DRAWINGS LIGHT UP AND RIGHT}' + HLIN        # └──
+
+
+def subclasses(cls):
+    try:
+        return cls.__subclasses__()
+    except TypeError:  # handle the `type` type
+        return cls.__subclasses__(cls)
+
+
+def tree(cls, level=0, last_sibling=True):
+    yield cls, level, last_sibling
+    children = subclasses(cls)
+    if children:
+        last = children[-1]
+        for child in children:
+            yield from tree(child, level+1, child is last)
 
 
 def render_lines(tree_iter):
@@ -13,8 +29,8 @@ def render_lines(tree_iter):
     prefix = ''
 
     for cls, level, last in tree_iter:
-        prefix = prefix[:4 * (level-1)]
-        prefix = prefix.replace(TEE, PIPE).replace(ELBOW, SP*4)
+        prefix = prefix[:4 * (level - 1)]
+        prefix = prefix.replace(TEE, VLIN).replace(ELBOW, SP * 4)
         prefix += ELBOW if last else TEE
         yield prefix + cls.__name__
 
