File size: 8,516 Bytes
481fbb6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 | #!/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",
]
|