mxpic_forge/generate_handbook.py

import ast
from dataclasses import dataclass
from pathlib import Path


DEFAULT_SRC_ROOT = Path("mxpic/components")
DEFAULT_DOCS_ROOT = Path("docs/source/mxpic/components")
DEFAULT_PACKAGE_ROOT = "mxpic.components"


@dataclass(frozen=True)
class ModuleDoc:
    """Documentation targets discovered in one Python module."""

    source_path: Path
    relative_path: Path
    module_name: str
    classes: tuple[str, ...]
    functions: tuple[str, ...]

    @property
    def output_path(self) -> Path:
        return self.relative_path.with_suffix(".md")


def public_top_level_members(py_file: Path) -> tuple[tuple[str, ...], tuple[str, ...]]:
    """Return public top-level classes and functions without importing a module."""
    tree = ast.parse(py_file.read_text(encoding="utf-8"), filename=str(py_file))
    classes: list[str] = []
    functions: list[str] = []
    seen_classes: set[str] = set()
    seen_functions: set[str] = set()

    for node in tree.body:
        if (
            isinstance(node, ast.ClassDef)
            and not node.name.startswith("_")
            and node.name not in seen_classes
        ):
            classes.append(node.name)
            seen_classes.add(node.name)
            continue

        if (
            isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
            and not node.name.startswith("_")
            and node.name not in seen_functions
        ):
            functions.append(node.name)
            seen_functions.add(node.name)

    return tuple(classes), tuple(functions)


def discover_modules(
    src_root: Path = DEFAULT_SRC_ROOT,
    package_root: str = DEFAULT_PACKAGE_ROOT,
) -> list[ModuleDoc]:
    """Discover Python modules that should receive generated Markdown pages."""
    src_root = Path(src_root)
    modules: list[ModuleDoc] = []

    for py_file in sorted(src_root.rglob("*.py")):
        if py_file.name == "__init__.py":
            continue

        relative_path = py_file.relative_to(src_root)
        module_suffix = ".".join(relative_path.with_suffix("").parts)
        classes, functions = public_top_level_members(py_file)
        modules.append(
            ModuleDoc(
                source_path=py_file,
                relative_path=relative_path,
                module_name=f"{package_root}.{module_suffix}",
                classes=classes,
                functions=functions,
            )
        )

    return modules


def remove_generated_markdown(docs_root: Path) -> None:
    """Remove generated Markdown while preserving images and other assets."""
    docs_root = Path(docs_root)
    if not docs_root.exists():
        return

    for md_file in docs_root.rglob("*.md"):
        md_file.unlink()


def write_module_page(module: ModuleDoc, docs_root: Path) -> None:
    """Write a Markdown page for one Python module."""
    md_file_path = Path(docs_root) / module.output_path
    md_file_path.parent.mkdir(parents=True, exist_ok=True)

    lines = [
        f"# {module.module_name}",
        "",
        "```{eval-rst}",
        f".. automodule:: {module.module_name}",
        "   :no-members:",
        "```",
        "",
    ]

    for class_name in module.classes:
        image_path = md_file_path.parent / f"{class_name}.png"
        lines.extend([f"## {class_name}", "", "```{eval-rst}"])

        if image_path.exists():
            lines.extend(
                [
                    f".. image:: {class_name}.png",
                    "   :align: center",
                    "   :width: 600px",
                    "",
                ]
            )

        lines.extend(
            [
                f".. autoclass:: {module.module_name}.{class_name}",
                "   :members:",
                "   :undoc-members:",
                "   :show-inheritance:",
                "```",
                "",
            ]
        )

    for function_name in module.functions:
        lines.extend(
            [
                f"## {function_name}",
                "",
                "```{eval-rst}",
                f".. autofunction:: {module.module_name}.{function_name}",
                "```",
                "",
            ]
        )

    md_file_path.write_text("\n".join(lines), encoding="utf-8")


def title_for_directory(relative_dir: Path) -> str:
    """Build a readable title for a generated directory index."""
    if relative_dir == Path("."):
        return "Components"

    return relative_dir.name.replace("_", " ").title()


def write_index_page(docs_root: Path, relative_dir: Path, entries: list[str]) -> None:
    """Write a MyST toctree index for one docs directory."""
    index_path = Path(docs_root) / relative_dir / "index.md"
    index_path.parent.mkdir(parents=True, exist_ok=True)

    lines = [
        f"# {title_for_directory(relative_dir)}",
        "",
        "```{toctree}",
        "   :maxdepth: 2",
        "",
    ]
    lines.extend(entries)
    lines.extend(["```", ""])
    index_path.write_text("\n".join(lines), encoding="utf-8")


def write_index_pages(modules: list[ModuleDoc], docs_root: Path) -> None:
    """Create recursive index.md files that include all generated module pages."""
    directories: set[Path] = {Path(".")}
    module_entries: dict[Path, list[str]] = {}

    for module in modules:
        directory = module.output_path.parent
        directories.add(directory)
        module_entries.setdefault(directory, []).append(module.output_path.stem)

        current = directory
        while current != Path("."):
            current = current.parent
            directories.add(current)

    for directory in sorted(directories, key=lambda path: (len(path.parts), path.as_posix()), reverse=True):
        child_dirs = sorted(
            child
            for child in directories
            if child != directory and child.parent == directory
        )
        entries = [f"{child.name}/index" for child in child_dirs]
        entries.extend(sorted(module_entries.get(directory, [])))

        if entries or directory == Path("."):
            write_index_page(docs_root, directory, entries)


def generate_markdown_handbook(
    src_root: Path = DEFAULT_SRC_ROOT,
    docs_root: Path = DEFAULT_DOCS_ROOT,
    package_root: str = DEFAULT_PACKAGE_ROOT,
) -> list[ModuleDoc]:
    """Generate Sphinx Markdown pages for mxpic component modules."""
    print("Starting mxPIC Markdown generation...")

    src_root = Path(src_root)
    docs_root = Path(docs_root)

    if not src_root.exists():
        raise FileNotFoundError(f"Source directory not found: {src_root}")

    docs_root.mkdir(parents=True, exist_ok=True)
    remove_generated_markdown(docs_root)

    modules = discover_modules(src_root=src_root, package_root=package_root)
    for module in modules:
        write_module_page(module, docs_root)
        print(f"Generated docs for: {module.module_name}")

    write_index_pages(modules, docs_root)

    print(f"Markdown generation complete. Updated {len(modules)} module pages.")
    return modules


if __name__ == "__main__":
    generate_markdown_handbook()