plot cyclic python imports

griffe.md

griffe_legend.md

import_graph.py

from pathlib import Path
import ast
from typing import Union

import networkx as nx
import matplotlib.pyplot as plt
import matplotlib.colors as clrs
import matplotlib.patches as ptc


class Node:  # read: module
    nodes: list["Node"] = list()  # sorted from long to short, so foo.bar.baz is found before foo.bar
    base: str = ""

    def __init__(self, name):
        self.name = name
        for idx in range(len(Node.nodes)):
            if Node.nodes[idx].name < self.name:
                Node.nodes.insert(idx, self)
                break
        else:
            # case: list empty or smallest new element
            Node.nodes.append(self)

        self.children: set["Node"] = {self}  # read: imports that happen in this module

    def add(self, other: Union["Node", str]):
        """Add a child to this Node.

        This method can only be called once all Nodes have been instantiated.

        Args:
            other: Node, either as an object or as a string, which will be evaluated as
                its longest match from all valid node names.
        """
        if isinstance(other, Node):
            self.children.add(other)
            return
        if isinstance(other, str):
            for node in Node.nodes:
                if other.startswith(node.name):
                    self.children.add(node)
                    return

    @classmethod
    def resolve_init_propagation(cls):
        """Make nodes inherit the imports of all parent __init__.py files.

        This method can only be called once all Nodes have had all their children added.
        """
        for node in cls.nodes:
            # seed node's imports to all direct sub-packages
            for other in cls.nodes:
                if other is node:
                    continue
                if node.name in other.name:  # found a sub-package
                    for child in node.children:
                        other.add(child)
        # reduce noise
        for node in cls.nodes:
            node.children.remove(node)


class ImportVisitor(ast.NodeVisitor):
    def __init__(self):
        self.imports: list[str] = []

    def visit_Import(self, node: ast.Import):
        for alias in node.names:
            self.imports.append(alias.name)

    def visit_ImportFrom(self, node: ast.ImportFrom):
        for alias in node.names:
            self.imports.append(f"{node.module}.{alias.name}")


def import_graph(package_base: str) -> dict[str, list[str]]:
    """Creates a directed graph of import statements.

    Args:
        package_base: Path to the source root, e.g. "/home/dev/my_lib/src/my_lib"

    Returns:
        Mapping of python module names to a list of all module names that it imports
    """
    node_data: dict[Node, list[str]] = {}
    package_name = Path(package_base).name
    
    for path in map(str, Path(package_base).rglob("*")):
        if not path.endswith(".py"):
            continue
        with open(path) as f:
            ast_ = ast.parse(f.read())
        visitor = ImportVisitor()
        visitor.visit(ast_)
        # how to turn a filename into a module path
        # TODO: handle namespaces and dist-name != folder-name
        name = path.removeprefix(package_base[:-len(package_name)]).removesuffix(".py").removesuffix("__init__").strip("/").replace("/", ".")
        node_data[Node(name)] = [e for e in visitor.imports if e.startswith(package_name)]
    
    for node, imports in node_data.items():
        for import_ in imports:
            node.add(import_)

    Node.resolve_init_propagation()
    for node in reversed(Node.nodes):
        print(node.name, [n.name for n in node.children])
    return {node.name: [n.name for n in node.children] for node in Node.nodes}


def draw_graph(graph: dict[str, list[str]], package: str):
    """Turns a graph dictionary into a picture."""
    all_colors = list(clrs.CSS4_COLORS)
    ratio = 0 if not graph else len(all_colors) / len(graph)
    colors = {k: all_colors[int(i*ratio)] for i, k in enumerate(graph)}

    g = nx.DiGraph()
    g.add_nodes_from(graph.keys())
    g.add_edges_from([(parent, child) for parent, children in graph.items() for child in children])
    circulars = [e for e in g.edges() if g.has_edge(*reversed(e))]
    layout = nx.kamada_kawai_layout(g)
    nx.draw_networkx_nodes(
        g,
        pos=layout,
        node_color=colors.values(),
        node_shape="o",
    )
    nx.draw_networkx_edges(
        g,
        pos=layout,
        arrows=True,
        edgelist=[e for e in g.edges() if e not in circulars],
    )
    nx.draw_networkx_edges(
        g,
        pos=layout,
        edge_color="red",
        arrows=True,
        edgelist=circulars,
    )
    plt.title(f"Import graph for {package}")

    plt.savefig(f"{package}.png")
    plt.show()
    plt.legend(
        loc=2,
        prop={"size": 6},
        handles=[ptc.Patch(color=c, label=l) for l, c in colors.items()]
        )
    plt.savefig(f"{package}_legend.png")


if __name__ == '__main__':
    draw_graph(import_graph("/home/me/dev/griffe/src/griffe"), "griffe")

output.txt

griffe ['griffe.loader', 'griffe.diff', 'griffe.git']
griffe.__main__ ['griffe.git', 'griffe.loader', 'griffe.cli', 'griffe.diff', 'griffe']
griffe.agents ['griffe.git', 'griffe.loader', 'griffe.diff', 'griffe']
griffe.agents.base ['griffe.agents', 'griffe.git', 'griffe.agents.nodes', 'griffe.loader', 'griffe.diff', 'griffe']
griffe.agents.extensions ['griffe.agents', 'griffe.git', 'griffe.loader', 'griffe.diff', 'griffe', 'griffe.extensions']
griffe.agents.inspector ['griffe.importer', 'griffe.agents', 'griffe.git', 'griffe.expressions', 'griffe.docstrings.parsers', 'griffe.agents.nodes', 'griffe.loader', 'griffe.agents.base', 'griffe.dataclasses', 'griffe.diff', 'griffe', 'griffe.extensions', 'griffe.collections']
griffe.agents.nodes ['griffe.agents', 'griffe.git', 'griffe.expressions', 'griffe.loader', 'griffe.exceptions', 'griffe.logger', 'griffe.dataclasses', 'griffe.diff', 'griffe', 'griffe.collections']
griffe.agents.visitor ['griffe.agents', 'griffe.git', 'griffe.expressions', 'griffe.docstrings.parsers', 'griffe.agents.nodes', 'griffe.loader', 'griffe.exceptions', 'griffe.agents.base', 'griffe.dataclasses', 'griffe.diff', 'griffe', 'griffe.extensions', 'griffe.collections']
griffe.cli ['griffe.extensions.base', 'griffe.git', 'griffe.docstrings.parsers', 'griffe.encoders', 'griffe.loader', 'griffe.exceptions', 'griffe.logger', 'griffe.diff', 'griffe', 'griffe.extensions']
griffe.collections ['griffe.git', 'griffe.loader', 'griffe.dataclasses', 'griffe.diff', 'griffe', 'griffe.mixins']
griffe.dataclasses ['griffe.git', 'griffe.expressions', 'griffe.docstrings.parsers', 'griffe.docstrings.dataclasses', 'griffe.loader', 'griffe.exceptions', 'griffe.diff', 'griffe', 'griffe.collections', 'griffe.mixins']
griffe.diff ['griffe.git', 'griffe.loader', 'griffe.logger', 'griffe.dataclasses', 'griffe']
griffe.docstrings ['griffe.git', 'griffe.docstrings.parsers', 'griffe.loader', 'griffe.diff', 'griffe']
griffe.docstrings.dataclasses ['griffe.docstrings', 'griffe.git', 'griffe.docstrings.parsers', 'griffe.loader', 'griffe.dataclasses', 'griffe.diff', 'griffe']
griffe.docstrings.google ['griffe.docstrings', 'griffe.diff', 'griffe.git', 'griffe.expressions', 'griffe.docstrings.parsers', 'griffe.docstrings.dataclasses', 'griffe.loader', 'griffe.dataclasses', 'griffe.docstrings.utils', 'griffe']
griffe.docstrings.markdown ['griffe.docstrings', 'griffe.git', 'griffe.docstrings.parsers', 'griffe.loader', 'griffe.diff', 'griffe']
griffe.docstrings.numpy ['griffe.docstrings', 'griffe.diff', 'griffe.git', 'griffe.expressions', 'griffe.docstrings.parsers', 'griffe.docstrings.dataclasses', 'griffe.loader', 'griffe.logger', 'griffe.dataclasses', 'griffe.docstrings.utils', 'griffe']
griffe.docstrings.parsers ['griffe.docstrings', 'griffe.git', 'griffe.docstrings.numpy', 'griffe.docstrings.dataclasses', 'griffe.loader', 'griffe.docstrings.google', 'griffe.dataclasses', 'griffe.docstrings.sphinx', 'griffe.diff', 'griffe']
griffe.docstrings.sphinx ['griffe.docstrings', 'griffe.diff', 'griffe.git', 'griffe.expressions', 'griffe.docstrings.parsers', 'griffe.docstrings.dataclasses', 'griffe.loader', 'griffe.dataclasses', 'griffe.docstrings.utils', 'griffe']
griffe.docstrings.utils ['griffe.docstrings', 'griffe.diff', 'griffe.git', 'griffe.expressions', 'griffe.docstrings.parsers', 'griffe.agents.nodes', 'griffe.loader', 'griffe.logger', 'griffe.dataclasses', 'griffe']
griffe.encoders ['griffe.diff', 'griffe.git', 'griffe.expressions', 'griffe.docstrings.parsers', 'griffe.docstrings.dataclasses', 'griffe.loader', 'griffe.dataclasses', 'griffe']
griffe.exceptions ['griffe.git', 'griffe.loader', 'griffe.dataclasses', 'griffe.diff', 'griffe']
griffe.expressions ['griffe.git', 'griffe.loader', 'griffe.exceptions', 'griffe.diff', 'griffe']
griffe.extensions ['griffe.extensions.base', 'griffe.git', 'griffe.loader', 'griffe.diff', 'griffe']
griffe.extensions.base ['griffe.extensions', 'griffe.git', 'griffe.agents.nodes', 'griffe.loader', 'griffe.exceptions', 'griffe.agents.base', 'griffe.agents.visitor', 'griffe.agents.inspector', 'griffe.diff', 'griffe', 'griffe.importer']
griffe.extensions.hybrid ['griffe.extensions', 'griffe.extensions.base', 'griffe.git', 'griffe.agents.nodes', 'griffe.loader', 'griffe.exceptions', 'griffe.logger', 'griffe.agents.visitor', 'griffe', 'griffe.diff', 'griffe.importer']
griffe.finder ['griffe.git', 'griffe.loader', 'griffe.exceptions', 'griffe.logger', 'griffe.dataclasses', 'griffe.diff', 'griffe']
griffe.git ['griffe.docstrings.parsers', 'griffe.loader', 'griffe.dataclasses', 'griffe.diff', 'griffe', 'griffe.extensions', 'griffe.collections']
griffe.importer ['griffe.git', 'griffe.loader', 'griffe.diff', 'griffe']
griffe.loader ['griffe.agents.inspector', 'griffe.finder', 'griffe.git', 'griffe.expressions', 'griffe.merger', 'griffe.docstrings.parsers', 'griffe.exceptions', 'griffe.logger', 'griffe.agents.visitor', 'griffe.dataclasses', 'griffe.stats', 'griffe.diff', 'griffe', 'griffe.extensions', 'griffe.collections']
griffe.logger ['griffe.git', 'griffe.loader', 'griffe.diff', 'griffe']
griffe.merger ['griffe.git', 'griffe.loader', 'griffe.exceptions', 'griffe.logger', 'griffe.dataclasses', 'griffe.diff', 'griffe']
griffe.mixins ['griffe.diff', 'griffe.git', 'griffe.merger', 'griffe.loader', 'griffe.exceptions', 'griffe.logger', 'griffe.dataclasses', 'griffe.encoders', 'griffe']
griffe.stats ['griffe.git', 'griffe.loader', 'griffe.exceptions', 'griffe.dataclasses', 'griffe.diff', 'griffe']

setup.sh

# for ubuntu
sudo apt-get install graphviz

python3.10 -m venv venv
./venv/bin/python -m pip install "networkx[default]" matplotlib pygraphviz pydot