Load and display Python class hierarchy

class_graph.py

"""
Build and display a graph of classes in a module.

Call from virtualenv of repo you're inspecting.

Example call, from inside edx-platform:

    DJANGO_SETTINGS_MODULE=lms.envs.test python3 ./class_graph.py xmodule.modulestore BulkOperationsMixin | dot -Tsvg > BulkOperationsMixin.svg
"""

import pkgutil
import re
import sys
from importlib import import_module


# Some nonsense to allow importing of modules in Django repos that try
# to load settings or work with models on module load.
try:
    import django
    django.setup()
except BaseException:
    pass


def main(in_module, *start_classes_rel):
    re_in = re.compile('^' + re.escape(in_module) + r'(\.|$)')
    def is_in_module(cls):
        return re_in.search(cls.__module__)

    # Make sure all submodules are loaded first (so that subclasses can be found)

    start_load = import_module(in_module)
    for m in pkgutil.walk_packages(start_load.__path__, start_load.__name__ + '.'):
        if re.search(r'(^|[_.])tests?([_.]|$)', m.name):
            continue
        import_module(m.name)

    # Load the initial set of classes

    classes = set()
    for crel in start_classes_rel:
        (full_mod, classname) = f'{in_module}.{crel}'.rsplit('.', 1)
        classes.add(getattr(import_module(full_mod), classname))

    # Transitive closure on class ancestors and descendents (within module)

    while True:
        wider = set()
        for cls in classes:
            wider.add(cls)
            wider |= set(cls.__subclasses__())
            wider |= set(cls.__bases__)
        wider = {c for c in wider if is_in_module(c)}
        expanded = wider == classes
        classes = wider
        if expanded:
            break

    # Print to DOT language for graphviz

    def outname(cls):
        """ Dotted module/class path relative to base module. """
        return (cls.__module__ + '.' + cls.__name__)[len(in_module):]

    print('digraph {')
    for cls in classes:
        for base in cls.__bases__:
            if is_in_module(base):
                # dot lays out top to bottom, but we want arrows to point upwards.
                # So, we have edges go down, but with a backwards arrow!
                print(f'  "{outname(base)}" -> "{outname(cls)}" [dir=back]')
    print('}')


if __name__ == '__main__':
    main(*sys.argv[1:])

example output - BulkOperationsMixin.svg