app/graph.py

import ast
from dataclasses import dataclass


@dataclass
class Graph:
    nodes: list[dict[str, str]]
    edges: list[dict[str, str]]


def _call_target_name(node: ast.AST) -> str | None:
    if isinstance(node, ast.Name):
        return node.id
    if isinstance(node, ast.Attribute) and isinstance(node.attr, str):
        return node.attr
    return None


def build_graph(code: str) -> Graph:
    tree = ast.parse(code)
    functions = {}
    classes: dict[str, list[str]] = {}

    for node in tree.body:
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            functions[node.name] = node
        elif isinstance(node, ast.ClassDef):
            method_names = []
            for child in node.body:
                if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)):
                    method_names.append(child.name)
            classes[node.name] = method_names

    nodes: list[dict[str, str]] = []
    edges: list[dict[str, str]] = []

    for class_name, method_names in classes.items():
        nodes.append({"id": class_name, "label": class_name, "type": "class"})
        for method_name in method_names:
            method_id = f"{class_name}.{method_name}"
            nodes.append({"id": method_id, "label": method_name, "type": "method"})
            edges.append({"source": class_name, "target": method_id, "type": "contains"})

    for func_name in functions:
        nodes.append({"id": func_name, "label": func_name, "type": "function"})

    known_nodes = {node["id"] for node in nodes}
    call_edges = set()

    def add_call_edge(source: str, target: str) -> None:
        if source == target:
            return
        if target not in known_nodes:
            return
        call_edges.add((source, target))

    for func_name, func_node in functions.items():
        for call in [n for n in ast.walk(func_node) if isinstance(n, ast.Call)]:
            target = _call_target_name(call.func)
            if target is None:
                continue
            if target in functions:
                add_call_edge(func_name, target)

    for class_name, method_names in classes.items():
        for node in tree.body:
            if isinstance(node, ast.ClassDef) and node.name == class_name:
                for child in node.body:
                    if not isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)):
                        continue
                    source_id = f"{class_name}.{child.name}"
                    for call in [n for n in ast.walk(child) if isinstance(n, ast.Call)]:
                        target = call.func
                        if isinstance(target, ast.Name):
                            if target.id in functions:
                                add_call_edge(source_id, target.id)
                            continue
                        if isinstance(target, ast.Attribute):
                            if isinstance(target.value, ast.Name):
                                if target.value.id == "self":
                                    target_id = f"{class_name}.{target.attr}"
                                    add_call_edge(source_id, target_id)
                                elif target.value.id in classes:
                                    target_id = f"{target.value.id}.{target.attr}"
                                    add_call_edge(source_id, target_id)

    for source, target in sorted(call_edges):
        edges.append({"source": source, "target": target, "type": "calls"})

    return Graph(nodes=nodes, edges=edges)