Create a MD document (and Mermaid diagram) of all the Python functions and their connections in a py file

functions_connections_docs_and_diagram.py

"""
analyse_functions.py
--------------------
Statically analyses a Python file and outputs:
  - A Markdown report of all functions (with signatures, docstrings, line numbers)
  - A Mermaid call-graph diagram (who calls whom)

Usage:
    python analyse_functions.py <target_file.py> [--out report.md]
"""

import ast
import sys
import argparse
from pathlib import Path
from dataclasses import dataclass, field
from collections import defaultdict


# ── Data structures ────────────────────────────────────────────────────────────

@dataclass
class FunctionInfo:
    name: str
    lineno: int
    end_lineno: int
    args: list[str]
    decorators: list[str]
    docstring: str | None
    calls: list[tuple[str, int]]   # (callee_name, line_number)
    is_method: bool = False
    class_name: str | None = None


# ── AST visitors ──────────────────────────────────────────────────────────────

class FunctionVisitor(ast.NodeVisitor):
    """Collects every function/method definition and its outgoing calls."""

    def __init__(self):
        self.functions: dict[str, FunctionInfo] = {}
        self._current_class: str | None = None

    # -- helpers ---------------------------------------------------------------

    def _arg_names(self, arguments: ast.arguments) -> list[str]:
        names = [a.arg for a in arguments.args]
        if arguments.vararg:
            names.append(f"*{arguments.vararg.arg}")
        names += [a.arg for a in arguments.kwonlyargs]
        if arguments.kwarg:
            names.append(f"**{arguments.kwarg.arg}")
        return names

    def _decorator_names(self, decorator_list) -> list[str]:
        names = []
        for d in decorator_list:
            if isinstance(d, ast.Name):
                names.append(d.id)
            elif isinstance(d, ast.Attribute):
                names.append(f"{ast.unparse(d)}")
            elif isinstance(d, ast.Call):
                names.append(f"{ast.unparse(d.func)}(...)")
        return names

    def _collect_calls(self, node: ast.FunctionDef) -> list[tuple[str, int]]:
        calls = []
        for child in ast.walk(node):
            if isinstance(child, ast.Call):
                if isinstance(child.func, ast.Name):
                    calls.append((child.func.id, child.lineno))
                elif isinstance(child.func, ast.Attribute):
                    calls.append((child.func.attr, child.lineno))
        return calls

    # -- visitors --------------------------------------------------------------

    def visit_ClassDef(self, node: ast.ClassDef):
        prev = self._current_class
        self._current_class = node.name
        self.generic_visit(node)
        self._current_class = prev

    def _visit_function(self, node):
        qualified = f"{self._current_class}.{node.name}" if self._current_class else node.name
        info = FunctionInfo(
            name=qualified,
            lineno=node.lineno,
            end_lineno=node.end_lineno,
            args=self._arg_names(node.args),
            decorators=self._decorator_names(node.decorator_list),
            docstring=ast.get_docstring(node),
            calls=self._collect_calls(node),
            is_method=self._current_class is not None,
            class_name=self._current_class,
        )
        self.functions[qualified] = info
        # visit nested functions
        self.generic_visit(node)

    visit_FunctionDef = _visit_function
    visit_AsyncFunctionDef = _visit_function


# ── Markdown generation ───────────────────────────────────────────────────────

def _signature(info: FunctionInfo) -> str:
    args = ", ".join(info.args)
    prefix = "async " if False else ""   # extend if you track async separately
    return f"`{info.name}({args})`"


def build_markdown(filepath: str, functions: dict[str, FunctionInfo]) -> str:
    path = Path(filepath)
    lines: list[str] = []

    lines.append(f"# Static Analysis — `{path.name}`\n")
    lines.append(f"> **File:** `{filepath}`  \n> **Functions found:** {len(functions)}\n")

    # ── Mermaid call-graph ─────────────────────────────────────────────────────
    lines.append("## Call Graph\n")
    lines.append("```mermaid")
    lines.append("flowchart TD")

    # Sanitise node IDs for Mermaid (no dots, brackets)
    def mid(name: str) -> str:
        return name.replace(".", "_").replace("<", "").replace(">", "")

    known = set(functions.keys())
    edges: set[tuple[str, str]] = set()

    for fname, info in functions.items():
        # node label
        label = fname.replace('"', "'")
        lines.append(f'    {mid(fname)}["{label}\\nL{info.lineno}–{info.end_lineno}"]')

    for fname, info in functions.items():
        for callee, lineno in info.calls:
            # resolve: try exact match, then suffix match
            resolved = None
            if callee in known:
                resolved = callee
            else:
                candidates = [k for k in known if k == callee or k.endswith(f".{callee}")]
                if candidates:
                    resolved = candidates[0]
            if resolved and resolved != fname:
                edge = (mid(fname), mid(resolved))
                if edge not in edges:
                    edges.add(edge)
                    lines.append(f"    {edge[0]} -->|L{lineno}| {edge[1]}")

    lines.append("```\n")

    # ── Function reference table ───────────────────────────────────────────────
    lines.append("## Function Reference\n")
    lines.append("| Function | Lines | Args | Decorators |")
    lines.append("|---|---|---|---|")
    for fname, info in sorted(functions.items(), key=lambda x: x[1].lineno):
        decs = ", ".join(f"`@{d}`" for d in info.decorators) or "—"
        args = ", ".join(f"`{a}`" for a in info.args) or "—"
        lines.append(
            f"| {_signature(info)} | {info.lineno}–{info.end_lineno} | {args} | {decs} |"
        )

    lines.append("")

    # ── Per-function detail ────────────────────────────────────────────────────
    lines.append("## Detailed Breakdown\n")

    for fname, info in sorted(functions.items(), key=lambda x: x[1].lineno):
        lines.append(f"### `{info.name}`\n")
        lines.append(f"- **Lines:** {info.lineno} – {info.end_lineno}  ")
        lines.append(f"- **File:** `{filepath}`  ")
        if info.class_name:
            lines.append(f"- **Class:** `{info.class_name}`  ")
        if info.decorators:
            lines.append(f"- **Decorators:** {', '.join(f'`@{d}`' for d in info.decorators)}  ")
        lines.append(f"- **Signature:** `{info.name}({', '.join(info.args)})`  ")

        if info.docstring:
            lines.append(f"\n> {info.docstring.splitlines()[0]}")

        # outgoing calls (only to known functions)
        known_calls = [
            (callee, lineno)
            for callee, lineno in info.calls
            if any(
                k == callee or k.endswith(f".{callee}")
                for k in functions
                if k != fname
            )
        ]
        if known_calls:
            lines.append("\n**Calls:**\n")
            seen = set()
            for callee, lineno in known_calls:
                key = (callee, lineno)
                if key not in seen:
                    seen.add(key)
                    lines.append(f"- → `{callee}` at line {lineno}")

        # incoming calls
        callers = [
            (n, ln)
            for n, inf in functions.items()
            for c, ln in inf.calls
            if (c == fname or fname.endswith(f".{c}")) and n != fname
        ]
        if callers:
            lines.append("\n**Called by:**\n")
            seen = set()
            for caller, lineno in callers:
                if (caller, lineno) not in seen:
                    seen.add((caller, lineno))
                    lines.append(f"- ← `{caller}` at line {lineno}")

        lines.append("")

    return "\n".join(lines)


# ── Entry point ───────────────────────────────────────────────────────────────

def analyse(filepath: str) -> dict[str, FunctionInfo]:
    source = Path(filepath).read_text(encoding="utf-8")
    tree = ast.parse(source, filename=filepath)
    visitor = FunctionVisitor()
    visitor.visit(tree)
    return visitor.functions


def main():
    parser = argparse.ArgumentParser(description="Static call-graph analyser for Python files.")
    parser.add_argument("file", help="Python file to analyse")
    parser.add_argument("--out", "-o", default=None, help="Output Markdown file (default: stdout)")
    args = parser.parse_args()

    if not Path(args.file).exists():
        print(f"Error: file not found: {args.file}", file=sys.stderr)
        sys.exit(1)

    functions = analyse(args.file)

    if not functions:
        print("No functions found.", file=sys.stderr)
        sys.exit(0)

    report = build_markdown(args.file, functions)

    if args.out:
        Path(args.out).write_text(report, encoding="utf-8")
        print(f"Report written to {args.out}")
    else:
        print(report)


if __name__ == "__main__":
    main()