stack/docs/scripts/extract_docs.py

"""Extract API docs from Python source using griffe.

Walks each package under ``src/``, parses NumPy-style docstrings via
griffe's pure-AST analysis (no imports needed), and writes one Markdown
file per module into ``docs/docs/api/{package}/{module}.md``.

Usage::

    uv run --with griffe python docs/scripts/extract_docs.py
"""

from __future__ import annotations

from pathlib import Path

from griffe import Class, Function, GriffeLoader, Module, Object

PACKAGES = ["aco", "api", "bcda", "bib", "bls", "ccw", "cms", "pfs", "rex"]
SRC_DIR = Path("src")
OUT_DIR = Path("docs/docs/api")


def _is_public(name: str) -> bool:
    return not name.startswith("_")


def _render_docstring(obj: Object) -> str:
    """Render an object's docstring as markdown."""
    if not obj.docstring:
        return ""
    return obj.docstring.value.strip() + "\n"


def _render_signature(func: Function) -> str:
    """Render a function signature."""
    params = []
    for p in func.parameters:
        if p.name in ("self", "cls"):
            continue
        part = p.name
        if p.annotation:
            part += f": {p.annotation}"
        if p.default is not None:
            default = str(p.default)
            # Truncate very long defaults
            if len(default) > 60:
                default = default[:57] + "..."
            part += f" = {default}"
        params.append(part)
    return f"({', '.join(params)})"


def _render_function(func: Function, heading: str = "###") -> str:
    """Render a function to markdown."""
    lines = [f"{heading} `{func.name}`\n"]
    sig = _render_signature(func)
    lines.append(f"```python\n{func.name}{sig}\n```\n")
    doc = _render_docstring(func)
    if doc:
        lines.append(doc)
    return "\n".join(lines) + "\n"


def _render_class(cls: Class) -> str:
    """Render a class and its public methods."""
    lines = [f"## `{cls.name}`\n"]
    doc = _render_docstring(cls)
    if doc:
        lines.append(doc)

    # __init__ signature
    if "__init__" in cls.members:
        init = cls.members["__init__"]
        if isinstance(init, Function):
            sig = _render_signature(init)
            lines.append(f"```python\n{cls.name}{sig}\n```\n")

    # Public methods
    methods = [
        m
        for name, m in cls.members.items()
        if isinstance(m, Function) and _is_public(name) and name != "__init__"
    ]
    if methods:
        lines.append("**Methods:**\n")
        for method in methods:
            lines.append(_render_function(method, heading="####"))

    return "\n".join(lines) + "\n"


def _render_module(mod: Module) -> str | None:
    """Render a module to markdown. Returns None if nothing public."""
    classes = [
        m for name, m in mod.members.items()
        if isinstance(m, Class) and _is_public(name)
    ]
    functions = [
        m for name, m in mod.members.items()
        if isinstance(m, Function) and _is_public(name)
    ]

    if not classes and not functions:
        return None

    # Module name without package prefix
    parts = mod.path.split(".")
    title = ".".join(parts)

    lines = [
        f"---\ntitle: {title}\n---\n",
        f"# `{title}`\n",
    ]

    doc = _render_docstring(mod)
    if doc:
        lines.append(doc)

    for cls in classes:
        lines.append(_render_class(cls))

    if functions:
        if classes:
            lines.append("---\n")
        lines.append("## Functions\n")
        for func in functions:
            lines.append(_render_function(func))

    return "\n".join(lines)


def extract_package(loader: GriffeLoader, package: str) -> int:
    """Extract docs for one package. Returns count of files written."""
    try:
        pkg = loader.load(package)
    except Exception as exc:
        print(f"  skip {package}: {exc}")
        return 0

    pkg_dir = OUT_DIR / package
    pkg_dir.mkdir(parents=True, exist_ok=True)

    # Write category metadata for Docusaurus sidebar
    cat_json = pkg_dir / "_category_.json"
    cat_json.write_text(
        f'{{"label": "{package}", "position": {PACKAGES.index(package) + 2}}}\n'
    )

    count = 0

    def _walk(mod: Module) -> None:
        nonlocal count
        content = _render_module(mod)
        if content:
            # Use module path relative to package for filename
            rel = mod.path.replace(f"{package}.", "").replace(".", "/")
            out_path = pkg_dir / f"{rel}.md"
            out_path.parent.mkdir(parents=True, exist_ok=True)
            out_path.write_text(content)
            count += 1
            print(f"  {mod.path} -> {out_path}")

        for name, member in mod.members.items():
            if isinstance(member, Module) and _is_public(name):
                _walk(member)

    _walk(pkg)
    return count


def main() -> None:
    # Clean output directory
    if OUT_DIR.exists():
        import shutil
        shutil.rmtree(OUT_DIR)
    OUT_DIR.mkdir(parents=True, exist_ok=True)

    loader = GriffeLoader(search_paths=[str(SRC_DIR)])
    total = 0

    for package in PACKAGES:
        print(f"Extracting {package}...")
        n = extract_package(loader, package)
        total += n

    print(f"\nDone: {total} module docs written to {OUT_DIR}")


if __name__ == "__main__":
    main()