File size: 7,585 Bytes
a39b25f | 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 | """Standalone Python function indexer for the CodeRAG demo."""
from __future__ import annotations
import ast
import hashlib
import json
import re
import textwrap
from pathlib import Path
from .records import FileIndex, FunctionRecord
SCHEMA_VERSION = "coderag-standalone-file-index-v1"
class _FunctionCollector(ast.NodeVisitor):
def __init__(self, source: str):
self.source = source
self.stack: list[tuple[str, str]] = []
self.functions: list[dict] = []
def visit_ClassDef(self, node: ast.ClassDef):
self.stack.append(("class", node.name))
self.generic_visit(node)
self.stack.pop()
def visit_FunctionDef(self, node: ast.FunctionDef):
self._visit_function(node)
def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef):
self._visit_function(node)
def _visit_function(self, node: ast.FunctionDef | ast.AsyncFunctionDef):
parents = [name for _, name in self.stack]
qualified_name = ".".join([*parents, node.name])
class_name = next(
(name for kind, name in reversed(self.stack) if kind == "class"),
None,
)
self.functions.append({
"node": node,
"name": node.name,
"qualified_name": qualified_name,
"class_name": class_name,
"symbol_type": "method" if class_name else "function",
})
self.stack.append(("function", node.name))
self.generic_visit(node)
self.stack.pop()
def _source_segment(source: str, node: ast.AST) -> str:
segment = ast.get_source_segment(source, node)
if segment:
return textwrap.dedent(segment).strip()
lines = source.splitlines()
start = getattr(node, "lineno", 1) - 1
end = getattr(node, "end_lineno", start + 1)
return "\n".join(lines[start:end]).strip()
def _signature(source: str, node: ast.FunctionDef | ast.AsyncFunctionDef) -> str:
first_line = _source_segment(source, node).splitlines()[0]
return first_line.rstrip(":")
def _function_id(file_path: str, qualified_name: str, start_line: int) -> str:
payload = f"{Path(file_path).as_posix()}::{qualified_name}::{start_line}"
return hashlib.sha1(payload.encode("utf-8")).hexdigest()[:16]
def _call_name(call: ast.Call) -> str | None:
func = call.func
if isinstance(func, ast.Name):
return func.id
if isinstance(func, ast.Attribute):
parts = [func.attr]
value = func.value
while isinstance(value, ast.Attribute):
parts.append(value.attr)
value = value.value
if isinstance(value, ast.Name):
parts.append(value.id)
return ".".join(reversed(parts))
return None
def _extract_calls(node: ast.AST) -> list[str]:
calls = []
for child in ast.walk(node):
if isinstance(child, ast.Call):
name = _call_name(child)
if name:
calls.append(name)
return list(dict.fromkeys(calls))
def _resolve_callees(calls: list[str], known_names: set[str]) -> tuple[list[str], list[str]]:
resolved = []
unresolved = []
short_to_qualified = {name.rsplit(".", 1)[-1]: name for name in known_names}
for call in calls:
if call in known_names:
resolved.append(call)
elif call.rsplit(".", 1)[-1] in short_to_qualified:
resolved.append(short_to_qualified[call.rsplit(".", 1)[-1]])
else:
unresolved.append(call)
return list(dict.fromkeys(resolved)), list(dict.fromkeys(unresolved))
def _words(value: str) -> str:
value = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", value)
value = value.replace("_", " ").replace(".", " ")
return re.sub(r"\s+", " ", value).strip().lower()
def _summary(name: str, node: ast.AST, docstring: str | None, calls: list[str]) -> str:
verbs = []
if any(isinstance(item, ast.Return) for item in ast.walk(node)):
verbs.append("returns a value")
if any(isinstance(item, (ast.For, ast.While, ast.AsyncFor)) for item in ast.walk(node)):
verbs.append("iterates over data")
if any(isinstance(item, ast.If) for item in ast.walk(node)):
verbs.append("branches on conditions")
if any(isinstance(item, ast.Raise) for item in ast.walk(node)):
verbs.append("raises errors")
pieces = [f"Function `{name}` handles {_words(name)}."]
if docstring:
pieces.append(f"Docstring: {docstring.strip().splitlines()[0]}")
if calls:
pieces.append("Calls: " + ", ".join(calls[:8]) + ".")
if verbs:
pieces.append("Behavior: " + ", ".join(verbs) + ".")
return " ".join(pieces)
def _propagated_text(record: FunctionRecord, by_name: dict[str, FunctionRecord]) -> str:
parts = [record.own_summary, record.raw_code]
for callee in record.direct_callees[:4]:
callee_record = by_name.get(callee)
if callee_record:
parts.append(f"Context from callee {callee}: {callee_record.own_summary}")
return "\n\n".join(part for part in parts if part)
def build_file_index(
file_path: str | Path,
*,
propagation_depth: int = 2,
max_callees: int = 4,
callee_context_tokens: int = 192,
) -> FileIndex:
del propagation_depth, max_callees, callee_context_tokens
path = Path(file_path)
source = path.read_text(encoding="utf-8")
tree = ast.parse(source)
collector = _FunctionCollector(source)
collector.visit(tree)
known_names = {item["qualified_name"] for item in collector.functions}
records = []
for item in collector.functions:
node = item["node"]
raw_code = _source_segment(source, node)
calls = _extract_calls(node)
direct_callees, unresolved_calls = _resolve_callees(calls, known_names)
docstring = ast.get_docstring(node)
own_summary = _summary(item["qualified_name"], node, docstring, calls)
record = FunctionRecord(
function_id=_function_id(str(path), item["qualified_name"], node.lineno),
name=item["name"],
qualified_name=item["qualified_name"],
symbol_type=item["symbol_type"],
file_path=str(path),
start_line=node.lineno,
end_line=getattr(node, "end_lineno", node.lineno),
signature=_signature(source, node),
docstring=docstring,
raw_code=raw_code,
own_sketch=f"{_signature(source, node)}\n{own_summary}",
own_summary=own_summary,
direct_callees=direct_callees,
unresolved_calls=unresolved_calls,
)
records.append(record)
by_name = {record.qualified_name: record for record in records}
for record in records:
record.callee_context = "\n".join(
by_name[callee].own_summary
for callee in record.direct_callees
if callee in by_name
)
record.propagated_text = _propagated_text(record, by_name)
return FileIndex(
schema_version=SCHEMA_VERSION,
file_path=str(path),
function_count=len(records),
call_edges=[
{"caller": record.qualified_name, "callee": callee}
for record in records
for callee in record.direct_callees
],
indexing_order=[record.qualified_name for record in records],
functions=records,
)
def write_file_index(index: FileIndex, output_path: str | Path) -> Path:
path = Path(output_path)
path.write_text(json.dumps(index.to_dict(), indent=2), encoding="utf-8")
return path
|