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