Aislo/docs/tools/extract_py.py

"""docs/tools/extract_py.py — 파이썬 추출(PLAN 478-1 · ast).

범위(`scope.json`)의 .py 를 읽어 `docs/_data/raw/functions_py.json` 에 씀(꼴 = PLAN 478 「추출 꼴」
+ FastAPI 라우트 데코레이터 `routes` [{method, path}]). 세는 규칙 = docs/tools/README.md.
"""

from __future__ import annotations

import ast
import json
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
SCOPE = json.loads((ROOT / "docs/tools/scope.json").read_text(encoding="utf-8"))
OUT = ROOT / "docs/_data/raw/functions_py.json"
HTTP = {"get", "post", "put", "delete", "patch", "head", "options"}


def rel(p: Path) -> str:
    return p.resolve().relative_to(ROOT).as_posix()


def skipped(r: str) -> bool:
    return any(s in r for s in SCOPE["skip"])


def module_file(mod: str) -> Path | None:
    base = ROOT.joinpath(*mod.split("."))
    for cand in (base.with_suffix(".py"), base / "__init__.py"):
        if cand.is_file():
            return cand
    return None


def resolve_from(file: Path, node: ast.ImportFrom) -> str:
    """`from <x> import` 의 절대 모듈 이름(상대 import 풀기)."""
    if not node.level:
        return node.module or ""
    parts = list(Path(rel(file)).parent.parts)
    parts = parts[: len(parts) - (node.level - 1)] if node.level > 1 else parts
    return ".".join([*parts, *([node.module] if node.module else [])])


def imports_of(file: Path, tree: ast.Module) -> dict[str, tuple[str, str | None]]:
    """이름 → (모듈, 그 모듈 안 이름 | None=모듈 자체)."""
    out: dict[str, tuple[str, str | None]] = {}
    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            for a in node.names:
                if a.asname:
                    out[a.asname] = (a.name, None)
                else:
                    out[a.name] = (a.name, None)
        elif isinstance(node, ast.ImportFrom):
            mod = resolve_from(file, node)
            for a in node.names:
                sub = f"{mod}.{a.name}" if mod else a.name
                out[a.asname or a.name] = (sub, None) if module_file(sub) else (mod, a.name)
    return out


def scope_files() -> list[Path]:
    seen: dict[str, Path] = {}
    queue = [p for r in SCOPE["roots"] for p in (ROOT / r).rglob("*.py")]
    # 계산 등록부(registry.json · extract_ts 가 먼저 씀)가 짚는 follow 폴더 파이썬도 범위
    registry = ROOT / "docs/_data/registry.json"
    named = json.loads(registry.read_text(encoding="utf-8")) if registry.is_file() else []
    for ref in (x for e in named for x in (e.get("pair"), e.get("runner")) if x):
        file = ref.removeprefix("py:").split("#")[0]
        if (
            file.endswith(".py")
            and any(file.startswith(f) for f in SCOPE["follow"])
            and (ROOT / file).is_file()
        ):
            queue.append(ROOT / file)
    while queue:
        f = queue.pop()
        r = rel(f)
        if r in seen or skipped(r):
            continue
        seen[r] = f
        tree = ast.parse(f.read_text(encoding="utf-8"), filename=r)
        for mod, _ in imports_of(f, tree).values():
            target = module_file(mod)
            if target and any(rel(target).startswith(x) for x in SCOPE["follow"]):
                queue.append(target)
    return [seen[k] for k in sorted(seen)]


def first_line(node: ast.AST) -> str | None:
    doc = (
        ast.get_docstring(node)
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
        else None
    )
    lines = [ln.strip() for ln in (doc or "").splitlines() if ln.strip()]
    return lines[0] if lines else None


def params_of(fn: ast.FunctionDef | ast.AsyncFunctionDef) -> list[dict]:
    a = fn.args
    pos = [*a.posonlyargs, *a.args]
    out = [
        {"name": p.arg, "type": ast.unparse(p.annotation) if p.annotation else None} for p in pos
    ]
    if a.vararg:
        out.append(
            {
                "name": f"*{a.vararg.arg}",
                "type": ast.unparse(a.vararg.annotation) if a.vararg.annotation else None,
            }
        )
    out += [
        {"name": p.arg, "type": ast.unparse(p.annotation) if p.annotation else None}
        for p in a.kwonlyargs
    ]
    if a.kwarg:
        out.append(
            {
                "name": f"**{a.kwarg.arg}",
                "type": ast.unparse(a.kwarg.annotation) if a.kwarg.annotation else None,
            }
        )
    return out


def router_prefixes(tree: ast.Module) -> dict[str, str]:
    """모듈 머리 `router = APIRouter(prefix="...")` → {router: prefix}."""
    out: dict[str, str] = {}
    for node in tree.body:
        if isinstance(node, ast.Assign) and isinstance(node.value, ast.Call):
            callee = node.value.func
            name = callee.id if isinstance(callee, ast.Name) else getattr(callee, "attr", "")
            if name != "APIRouter":
                continue
            prefix = next(
                (
                    k.value.value
                    for k in node.value.keywords
                    if k.arg == "prefix" and isinstance(k.value, ast.Constant)
                ),
                "",
            )
            for t in node.targets:
                if isinstance(t, ast.Name):
                    out[t.id] = prefix
    return out


def routes_of(fn: ast.FunctionDef | ast.AsyncFunctionDef, prefixes: dict[str, str]) -> list[dict]:
    out = []
    for d in fn.decorator_list:
        if not (isinstance(d, ast.Call) and isinstance(d.func, ast.Attribute)):
            continue
        verb = d.func.attr
        owner = d.func.value.id if isinstance(d.func.value, ast.Name) else ""
        path = d.args[0].value if d.args and isinstance(d.args[0], ast.Constant) else None
        if path is None or (verb not in HTTP and verb != "api_route"):
            continue
        methods = [verb.upper()] if verb in HTTP else ["GET"]
        for k in d.keywords:
            if (
                verb == "api_route"
                and k.arg == "methods"
                and isinstance(k.value, (ast.List, ast.Tuple))
            ):
                methods = [e.value.upper() for e in k.value.elts if isinstance(e, ast.Constant)]
        out += [{"method": m, "path": prefixes.get(owner, "") + path} for m in methods]
    return out


def own_calls(fn: ast.AST):
    """fn 몸통의 부름 — 안쪽 def · class 몫은 빼고 lambda · 내포는 fn 몫."""
    for child in ast.iter_child_nodes(fn):
        if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
            continue
        if isinstance(child, ast.Call):
            yield child
        yield from own_calls(child)


class Collector(ast.NodeVisitor):
    def __init__(self, file: Path, tree: ast.Module, used: dict[str, int]):
        self.file = rel(file)
        self.imports = imports_of(file, tree)
        self.prefixes = router_prefixes(tree)
        self.used = used
        self.stack: list[tuple[str, str]] = []  # (이름, "class" | "def")
        self.records: list[dict] = []
        self.raw_calls: dict[str, list[tuple[ast.expr, list[str], str | None]]] = {}

    def _qual(self, name: str) -> str:
        return ".".join([*(n for n, _ in self.stack), name])

    def visit_ClassDef(self, node: ast.ClassDef) -> None:
        self.stack.append((node.name, "class"))
        self.generic_visit(node)
        self.stack.pop()

    def _fn(self, node: ast.FunctionDef | ast.AsyncFunctionDef) -> None:
        qual = self._qual(node.name)
        rid = f"py:{self.file}#{qual}"
        self.used[rid] = self.used.get(rid, 0) + 1
        if self.used[rid] > 1:
            rid = f"{rid}~{self.used[rid]}"
        in_class = bool(self.stack) and self.stack[-1][1] == "class"
        top = not self.stack or (in_class and len(self.stack) == 1)
        rec = {
            "id": rid,
            "lang": "py",
            "file": self.file,
            "line": node.lineno,
            "name": qual,
            "kind": "method" if in_class else "function",
            "exported": top and not node.name.startswith("_"),
            "doc": first_line(node),
            "params": params_of(node),
            "returns": ast.unparse(node.returns) if node.returns else None,
            "calls": [],
            "called_by": [],
        }
        routes = routes_of(node, self.prefixes)
        if routes:
            rec["routes"] = routes
        self.records.append(rec)
        cls = self.stack[-1][0] if in_class else None
        # 이름 찾는 차례 = 자기 안 → 감싼 def(안쪽부터) → 모듈(클래스 칸은 건너뜀 = 파이썬 규칙)
        defs = [
            ".".join(n for n, _ in self.stack[: i + 1])
            for i, (_, k) in enumerate(self.stack)
            if k == "def"
        ]
        scopes = [qual, *reversed(defs), ""]
        self.raw_calls[rid] = [(c.func, scopes, cls) for c in own_calls(node)]
        self.stack.append((node.name, "def"))
        self.generic_visit(node)
        self.stack.pop()

    visit_FunctionDef = _fn
    visit_AsyncFunctionDef = _fn

    def resolve(
        self, func: ast.expr, scopes: list[str], cls: str | None, ids: set[str]
    ) -> str | None:
        def local(qual: str) -> str | None:
            rid = f"py:{self.file}#{qual}"
            return rid if rid in ids else None

        def external(mod: str, name: str) -> str | None:
            target = module_file(mod)
            rid = f"py:{rel(target)}#{name}" if target else None
            return rid if rid in ids else None

        if isinstance(func, ast.Name):
            for s in scopes:
                hit = local(f"{s}.{func.id}" if s else func.id)
                if hit:
                    return hit
            if func.id in self.imports:
                mod, name = self.imports[func.id]
                return external(mod, name) if name else None
            return None
        if isinstance(func, ast.Attribute) and isinstance(func.value, ast.Name):
            base = func.value.id
            if base in ("self", "cls") and cls:
                return local(f"{cls}.{func.attr}")
            if local(f"{base}.{func.attr}"):
                return local(f"{base}.{func.attr}")
            if base in self.imports:
                mod, name = self.imports[base]
                return (
                    external(mod, func.attr)
                    if name is None
                    else external(mod, f"{name}.{func.attr}")
                )
        return None


def main() -> None:
    used: dict[str, int] = {}
    collectors = []
    for f in scope_files():
        tree = ast.parse(f.read_text(encoding="utf-8"), filename=rel(f))
        c = Collector(f, tree, used)
        c.visit(tree)
        collectors.append(c)
    records = [r for c in collectors for r in c.records]
    ids = {r["id"] for r in records}
    index = {r["id"]: r for r in records}
    for c in collectors:
        for rid, calls in c.raw_calls.items():
            hits = {c.resolve(func, scopes, cls, ids) for func, scopes, cls in calls}
            index[rid]["calls"] = sorted(h for h in hits if h and h != rid)
    for r in records:
        for target in r["calls"]:
            index[target]["called_by"].append(r["id"])
    for r in records:
        r["called_by"] = sorted(set(r["called_by"]))
    records.sort(key=lambda r: r["id"])
    OUT.parent.mkdir(parents=True, exist_ok=True)
    OUT.write_text(json.dumps(records, ensure_ascii=False, indent=1) + "\n", encoding="utf-8")
    routes = sum(len(r.get("routes", [])) for r in records)
    print(f"py: 파일 {len(collectors)} · 함수 {len(records)} · 라우트 데코레이터 {routes}")


if __name__ == "__main__":
    main()