File size: 2,279 Bytes
71d239c | 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 | """Extract structure from a Python file using the standard-library `ast` module.
We pull out the facts a story needs — top-level functions and classes, their
first docstring line, and every import target — without running the code. This
is what lets a small model narrate reliably: it never has to *infer* structure.
"""
from __future__ import annotations
import ast
from schema import Symbol
def parse(text: str) -> tuple[list[Symbol], list[str]]:
"""Return (symbols, import_targets). Best-effort; never raises."""
try:
tree = ast.parse(text)
except SyntaxError:
return [], []
symbols: list[Symbol] = []
imports: list[str] = []
for node in ast.iter_child_nodes(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
symbols.append(Symbol(kind="function", name=node.name,
lineno=node.lineno, doc=_doc(node)))
elif isinstance(node, ast.ClassDef):
symbols.append(Symbol(kind="class", name=node.name,
lineno=node.lineno, doc=_doc(node)))
# Imports can be nested (inside functions/try), so walk the whole tree.
# For `from pkg import a, b` we emit candidates "pkg", "pkg.a", "pkg.b" so the
# graph resolver can link to pkg/a.py even when pkg is a package — longest
# match wins there.
for node in ast.walk(tree):
if isinstance(node, ast.Import):
imports.extend(alias.name for alias in node.names)
elif isinstance(node, ast.ImportFrom):
mod = node.module or ""
prefix = "." * (node.level or 0) # preserve relative-import depth
base = prefix + mod
imports.append(base or prefix)
for alias in node.names:
if alias.name and alias.name != "*":
imports.append((base + "." if base and not base.endswith(".") else base)
+ alias.name)
# De-dup, keep order.
seen: set[str] = set()
uniq = [i for i in imports if not (i in seen or seen.add(i))]
return symbols, uniq
def _doc(node: ast.AST) -> str | None:
doc = ast.get_docstring(node)
if not doc:
return None
return doc.strip().splitlines()[0][:140]
|