feat/stubgen: add struct and enum generator

src/capnp_stubgen/plugin.py

@@ -0,0 +1,123 @@
+"""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_enum import generate_enum
+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)
+
+    # ── 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)
+
+    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)
+
+    return emitter.render()