memory_lora/codegraph.py

#!/usr/bin/env python3
"""AST-based codebase graph representation.

New data representation beyond what the Code2LoRA paper used (raw file
chunks, mean+max pooled). This module extracts a STRUCTURAL summary of a
Python codebase -- imports (dependency edges between files), class
hierarchies, function/method signatures, and a best-effort call graph --
and serializes it into compact text sections. Those sections are fed
through the SAME frozen embedding pipeline as raw code
(``memory_lora.encoder.embed_document`` already accepts a list of
``(section_name, section_text)`` tuples), so a repository can be embedded
from BOTH its raw source text and its structural graph, without any
change to the encoder or hypernetwork.

Why this matters for "recall codebases": raw-text chunking captures
surface content (docstrings, comments, literal code) but dilutes the
signal that actually matters for API-level questions -- "what does
function X take", "what inherits from Y", "who imports Z". A compact
graph serialization puts that signal in a small number of dense tokens
instead of burying it across thousands of raw-code tokens.
"""

from __future__ import annotations

import ast
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Optional, Tuple


@dataclass
class FunctionSig:
    name: str
    args: List[str]
    returns: Optional[str]
    is_method: bool = False
    decorators: List[str] = field(default_factory=list)
    calls: List[str] = field(default_factory=list)  # best-effort, unresolved names


@dataclass
class ClassSig:
    name: str
    bases: List[str]
    methods: List[FunctionSig] = field(default_factory=list)


@dataclass
class FileGraph:
    path: str
    imports: List[str]
    import_froms: List[Tuple[str, List[str]]]  # (module, [names])
    classes: List[ClassSig]
    functions: List[FunctionSig]  # module-level only


def _annotation_to_str(node: Optional[ast.AST]) -> Optional[str]:
    if node is None:
        return None
    try:
        return ast.unparse(node)
    except Exception:  # noqa: BLE001
        return None


def _arg_to_str(a: ast.arg) -> str:
    ann = _annotation_to_str(a.annotation)
    return f"{a.arg}: {ann}" if ann else a.arg


def _extract_calls(node: ast.AST) -> List[str]:
    calls = []
    for n in ast.walk(node):
        if isinstance(n, ast.Call):
            f = n.func
            if isinstance(f, ast.Name):
                calls.append(f.id)
            elif isinstance(f, ast.Attribute):
                calls.append(f.attr)
    # dedupe, keep order
    seen = set()
    out = []
    for c in calls:
        if c not in seen:
            seen.add(c)
            out.append(c)
    return out[:20]  # cap -- this is a signal, not a full trace


def _function_sig(node, is_method: bool = False) -> FunctionSig:
    args = [_arg_to_str(a) for a in node.args.args]
    returns = _annotation_to_str(node.returns)
    decorators = [_annotation_to_str(d) or "" for d in node.decorator_list]
    return FunctionSig(
        name=node.name, args=args, returns=returns, is_method=is_method,
        decorators=[d for d in decorators if d], calls=_extract_calls(node),
    )


def extract_file_graph(source: str, path: str) -> Optional[FileGraph]:
    """Parse one Python file's source into a :class:`FileGraph`.
    Returns None on a syntax error (skip the file, don't crash the repo)."""
    try:
        tree = ast.parse(source)
    except (SyntaxError, ValueError):
        return None

    imports: List[str] = []
    import_froms: List[Tuple[str, List[str]]] = []
    classes: List[ClassSig] = []
    functions: List[FunctionSig] = []

    for node in ast.iter_child_nodes(tree):
        if isinstance(node, ast.Import):
            imports.extend(a.name for a in node.names)
        elif isinstance(node, ast.ImportFrom):
            mod = node.module or ("." * node.level)
            import_froms.append((mod, [a.name for a in node.names]))
        elif isinstance(node, ast.ClassDef):
            bases = [_annotation_to_str(b) or "?" for b in node.bases]
            methods = [
                _function_sig(n, is_method=True)
                for n in node.body
                if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef))
            ]
            classes.append(ClassSig(name=node.name, bases=bases, methods=methods))
        elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            functions.append(_function_sig(node, is_method=False))

    return FileGraph(path=path, imports=imports, import_froms=import_froms,
                      classes=classes, functions=functions)


def serialize_file_graph(g: FileGraph) -> str:
    """Compact text serialization -- dense signature summary, not prose."""
    lines = [f"# {g.path}"]
    if g.imports:
        lines.append("imports: " + ", ".join(g.imports))
    for mod, names in g.import_froms:
        lines.append(f"from {mod} import " + ", ".join(names))
    for c in g.classes:
        base_str = f"({', '.join(c.bases)})" if c.bases else ""
        lines.append(f"class {c.name}{base_str}:")
        for m in c.methods:
            dec = "".join(f"@{d} " for d in m.decorators)
            args = ", ".join(m.args)
            ret = f" -> {m.returns}" if m.returns else ""
            calls = f"  # calls: {', '.join(m.calls[:6])}" if m.calls else ""
            lines.append(f"  {dec}def {m.name}({args}){ret}{calls}")
    for fn in g.functions:
        dec = "".join(f"@{d} " for d in fn.decorators)
        args = ", ".join(fn.args)
        ret = f" -> {fn.returns}" if fn.returns else ""
        calls = f"  # calls: {', '.join(fn.calls[:6])}" if fn.calls else ""
        lines.append(f"{dec}def {fn.name}({args}){ret}{calls}")
    return "\n".join(lines)


def extract_repo_graph_sections(
    repo_dir: Path, max_files: int = 200, skip_dirs: Optional[set] = None,
) -> List[Tuple[str, str]]:
    """Walk a repo, extract + serialize each .py file's structural graph.
    Returns ``[(section_name, section_text), ...]`` -- directly usable as
    the ``sections`` argument to ``memory_lora.encoder.embed_document``,
    alongside (or instead of) raw-text sections.
    """
    skip_dirs = skip_dirs or {".git", "__pycache__", ".venv", "venv", "node_modules",
                               "build", "dist", ".tox", ".mypy_cache"}
    sections: List[Tuple[str, str]] = []
    n = 0
    for path in sorted(repo_dir.rglob("*.py")):
        if any(part in skip_dirs for part in path.parts):
            continue
        try:
            source = path.read_text(encoding="utf-8", errors="ignore")
        except OSError:
            continue
        rel = str(path.relative_to(repo_dir))
        g = extract_file_graph(source, rel)
        if g is None:
            continue
        if not (g.imports or g.import_froms or g.classes or g.functions):
            continue  # empty file, no structural signal
        sections.append((f"graph:{rel}", serialize_file_graph(g)))
        n += 1
        if n >= max_files:
            break
    return sections


def extract_repo_dependency_summary(repo_dir: Path, skip_dirs: Optional[set] = None) -> str:
    """One dense paragraph: which files import which other in-repo modules.
    A cheap approximation of a dependency graph edge list, useful as a
    single extra section capturing repo-wide (not per-file) structure."""
    skip_dirs = skip_dirs or {".git", "__pycache__", ".venv", "venv", "node_modules"}
    edges: List[str] = []
    module_names = set()
    files = [p for p in repo_dir.rglob("*.py") if not any(part in skip_dirs for part in p.parts)]
    for p in files:
        mod = str(p.relative_to(repo_dir)).replace("/", ".").removesuffix(".py")
        module_names.add(mod)

    for p in files:
        try:
            source = p.read_text(encoding="utf-8", errors="ignore")
            tree = ast.parse(source)
        except (SyntaxError, ValueError, OSError):
            continue
        rel = str(p.relative_to(repo_dir))
        for node in ast.walk(tree):
            if isinstance(node, ast.ImportFrom) and node.module:
                if node.module in module_names or any(m.startswith(node.module + ".") for m in module_names):
                    edges.append(f"{rel} -> {node.module}")
    return "dependency edges:\n" + "\n".join(edges[:300]) if edges else ""


__all__ = [
    "FunctionSig", "ClassSig", "FileGraph",
    "extract_file_graph", "serialize_file_graph",
    "extract_repo_graph_sections", "extract_repo_dependency_summary",
]