Generate a graph from class dependencies

graph.py

"""Create a graph of dependencies

Depends on graphviz being installed:

`brew/apt-get install graphviz`

`pip install graphviz`
"""
import argparse
import inspect
from typing import Type

import graphviz
import importlib


def import_class(import_path: str) -> Type:
    # Split the import path into module and class components
    module_name, class_name = import_path.rsplit(".", 1)

    # Import the module dynamically
    module = importlib.import_module(module_name)

    # Access the class from the module
    return getattr(module, class_name)


def extract_dependencies(klass: Type) -> dict[str, Type]:
    # Get the constructor (__init__) method of the class
    constructor = getattr(klass, "__init__")

    if not constructor:
        raise StopIteration("type without constructor")

    # Get the constructor's parameters and their types
    parameters = inspect.signature(constructor).parameters

    # Extract argument names and types
    arg_types = {}
    for param_name, param in parameters.items():
        if param.annotation != inspect.Parameter.empty:
            arg_types[param_name] = param.annotation

    return arg_types


def link_classes(current_class: Type, _dot: graphviz.Digraph) -> None:
    try:
        arg_types = extract_dependencies(current_class)
    except StopIteration:
        return

    for child_class in arg_types.values():
        # Add node (child class)
        _dot.node(str(child_class), shape="box", style="filled")
        # Add edges (dependencies)
        _dot.edge(str(current_class), str(child_class))

        link_classes(child_class, _dot)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="build dependency graph")
    parser.add_argument("class_path", type=str, help="Path to the class")
    parser.add_argument("--output", type=str, default="dependency_graph.svg", help="Output file")

    args = parser.parse_args()

    *names, extension = args.output.split(".")

    # Create a new Graphviz graph
    dot = graphviz.Digraph(format=extension)

    klass = import_class(args.class_path)
    dot.node(str(klass), shape="box", style="filled")

    link_classes(klass, dot)

    # Save the graph to a file (e.g., dependency_graph.png)
    dot.render(outfile=args.output, view=True)