analyzer/python_ast.py

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