capnpc-py/src/capnp_stubgen/plugin.py

"""Cap'n Proto compiler plugin entry point.

This is the ``capnpc-py`` binary. When invoked by ``capnp compile -opy``,
it reads a ``CodeGeneratorRequest`` from stdin and writes ``.pyi`` type
stub files to the current working directory.

Uses ``.to_dict()`` on the request to avoid C++ API quirks with
DynamicStruct-wrapped schema nodes (same approach as pycapnp's own
``_gen.py``).
"""

from __future__ import annotations

import os
import sys
from typing import TYPE_CHECKING

from .emitter import Emitter
from .gen_const import generate_const
from .gen_enum import generate_enum
from .gen_interface import fixup_interface_methods, generate_interface
from .gen_struct import generate_struct
from .models import TypeRegistry
from .schema_walker import build_type_registry
from .utils import filename_to_module_name, type_id_to_int

if TYPE_CHECKING:
    from typing import Any


def main() -> int:
    """Entry point. Returns 0 on success, 1 on error."""
    try:
        return _main_impl()
    except Exception as exc:
        print(f"capnpc-py: error: {exc}", file=sys.stderr)
        return 1


def _main_impl() -> int:
    import capnp

    # ── 1. Read CodeGeneratorRequest from stdin ────────────────────────
    request = _read_request()

    # Convert to dict for processing (avoids C++ DynamicStruct quirks)
    request_dict = request.to_dict()
    nodes: list[dict] = request_dict["nodes"]
    requested_files: list[dict] = request_dict["requestedFiles"]

    # ── 2. Build type registry ─────────────────────────────────────────
    registry = build_type_registry(nodes, requested_files)

    # Attach method param/result structs to their interfaces
    fixup_interface_methods(registry)

    # Record file ID → module name mappings
    # Register both requested files and their imports
    for req_file in requested_files:
        file_id = type_id_to_int(req_file["id"])
        module_name = filename_to_module_name(req_file["filename"])
        registry.set_file_module_name(file_id, module_name)

        # Also register imported files
        for imp in req_file.get("imports", []):
            imp_id = type_id_to_int(imp["id"])
            imp_name = filename_to_module_name(imp["name"])
            registry.set_file_module_name(imp_id, imp_name)

    # ── 3. Generate .pyi for each requested file ───────────────────────
    ok = True
    for req_file in requested_files:
        filename: str = req_file["filename"]
        file_id = type_id_to_int(req_file["id"])

        try:
            content = _generate_file(file_id, filename, registry)
            module_name = filename_to_module_name(filename)
            output_path = f"{module_name}.pyi"

            with open(output_path, "w") as f:
                f.write(content)

        except Exception as exc:
            print(
                f"capnpc-py: error generating stubs for {filename}: {exc}",
                file=sys.stderr,
            )
            ok = False

    return 0 if ok else 1


# ── request reading ────────────────────────────────────────────────────────


def _read_request() -> Any:
    """Read ``CodeGeneratorRequest`` from stdin.

    Uses pycapnp's ``SchemaParser`` to dynamically load ``schema.capnp``
    (shipped with pycapnp), then casts the stdin message.
    """
    import capnp

    capnp_dir = os.path.dirname(capnp.__file__)
    schema_path = os.path.join(capnp_dir, "schema.capnp")
    import_base = os.path.dirname(capnp_dir)

    parser = capnp.SchemaParser()
    schema_module = parser.load(schema_path, imports=[import_base])

    return schema_module.CodeGeneratorRequest.read(sys.stdin.buffer)


# ── file generation ────────────────────────────────────────────────────────


def _generate_file(
    file_id: int,
    filename: str,
    registry: TypeRegistry,
) -> str:
    """Generate the complete .pyi content for one schema file."""
    emitter = Emitter(source_filename=filename)

    # Collect cross-file imports
    _resolve_cross_imports(file_id, filename, registry, emitter)

    top_level = registry.get_top_level_types(file_id)

    _ORDER = {"enum": 0, "struct": 1, "const": 2, "interface": 3}

    def _sort_key(info):
        return (_ORDER.get(info.kind, 99), info.name)

    top_level.sort(key=_sort_key)

    for type_info in top_level:
        if type_info.kind == "enum":
            generate_enum(emitter, type_info)
        elif type_info.kind == "struct":
            generate_struct(emitter, type_info, registry)
        elif type_info.kind == "interface":
            generate_interface(emitter, type_info, registry)
        elif type_info.kind == "const":
            generate_const(emitter, type_info, registry)

    return emitter.render()


def _resolve_cross_imports(
    file_id: int,
    filename: str,
    registry: TypeRegistry,
    emitter: Emitter,
) -> None:
    """Emit Python import lines for types referenced from other files.

    Walks all types in *file_id*, collects field type references that
    belong to a different file, and emits relative import lines.
    """
    external_types: dict[str, set[str]] = {}  # module_name → {TypeName, TypeNameBuilder, ...}

    # Collect all types belonging to this file
    own_types: set[int] = set()
    for info in registry._types.values():
        if info.file_id == file_id and info.kind in ("struct", "enum", "interface"):
            own_types.add(info.type_id)

    # Walk types and find external references
    for info in registry._types.values():
        if info.file_id != file_id:
            continue

        refs = _collect_type_references(info, registry)
        for ref_id in refs:
            if ref_id not in own_types:
                ref_info = registry.get(ref_id)
                if ref_info is not None and ref_info.file_id is not None:
                    ref_file_id = ref_info.file_id
                    # Get the module name for the referenced file
                    module_name = _file_id_to_module_name(ref_file_id, registry)
                    if module_name:
                        if module_name not in external_types:
                            external_types[module_name] = set()
                        external_types[module_name].add(ref_info.scoped_name)

    # Emit import lines
    for module_name, type_names in sorted(external_types.items()):
        # Only emit relative imports (not self-imports)
        own_module = filename_to_module_name(filename)
        if module_name == own_module:
            continue

        # Collect all import names for this module
        all_imports: list[str] = []
        for tname in sorted(type_names):
            ref_info = None
            for info in registry._types.values():
                if info.scoped_name == tname and info.file_id != file_id:
                    ref_info = info
                    break

            all_imports.append(tname)
            if ref_info and ref_info.kind == "struct":
                all_imports.extend([f"{tname}Builder", f"{tname}Reader"])

        emitter.add_import(
            f"from .{module_name} import {', '.join(all_imports)}"
        )


def _collect_type_references(info: TypeInfo, registry: TypeRegistry) -> set[int]:
    """Collect all type IDs referenced by a TypeInfo's fields."""
    refs: set[int] = set()
    node = info.node

    if info.kind == "struct":
        for field in node.get("struct", {}).get("fields", []):
            _collect_field_references(field, refs, registry)
    elif info.kind == "interface":
        for method in node.get("interface", {}).get("methods", []):
            refs.add(type_id_to_int(method["paramStructType"]))
            refs.add(type_id_to_int(method["resultStructType"]))

    return refs


def _collect_field_references(field: dict, refs: set[int], registry: TypeRegistry) -> None:
    """Recursively collect type IDs from a field's type."""
    if "slot" in field:
        _collect_type_references_from_dict(field["slot"].get("type", {}), refs)
    elif "group" in field:
        refs.add(type_id_to_int(field["group"]["typeId"]))


def _collect_type_references_from_dict(type_dict: dict, refs: set[int]) -> None:
    """Recursively collect type IDs from a type dict."""
    if "struct" in type_dict:
        refs.add(type_id_to_int(type_dict["struct"]["typeId"]))
    elif "enum" in type_dict:
        refs.add(type_id_to_int(type_dict["enum"]["typeId"]))
    elif "interface" in type_dict:
        refs.add(type_id_to_int(type_dict["interface"]["typeId"]))
    elif "list" in type_dict:
        _collect_type_references_from_dict(
            type_dict["list"].get("elementType", {}), refs
        )


def _file_id_to_module_name(file_id: int, registry: TypeRegistry) -> str | None:
    """Convert a file node ID to its Python module name."""
    return registry.get_file_module_name(file_id)