| """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 |
|
|
