From 5ba32908b44283571bd2346d74bd6d81812d5ca8 Mon Sep 17 00:00:00 2001 From: worldmozara Date: Thu, 9 Jul 2026 15:39:10 +0800 Subject: [PATCH] feat/stubgen: add interface support --- src/capnp_stubgen/emitter.py | 10 +- src/capnp_stubgen/gen_interface.py | 150 +++++++++++++++++++++++++++++ src/capnp_stubgen/gen_struct.py | 2 +- src/capnp_stubgen/models.py | 14 +++ src/capnp_stubgen/plugin.py | 127 ++++++++++++++++++++++++ 5 files changed, 300 insertions(+), 3 deletions(-) create mode 100644 src/capnp_stubgen/gen_interface.py diff --git a/src/capnp_stubgen/emitter.py b/src/capnp_stubgen/emitter.py index bf9f02e..1c40191 100644 --- a/src/capnp_stubgen/emitter.py +++ b/src/capnp_stubgen/emitter.py @@ -162,9 +162,15 @@ class Emitter: # Blank line between header and imports parts.append("") - # Regular imports + # Regular imports (local `.` imports last, after stdlib/third-party) if self._imports: - for imp in sorted(self._imports): + local_imports = sorted(i for i in self._imports if i.startswith("from .")) + std_imports = sorted(i for i in self._imports if not i.startswith("from .")) + for imp in std_imports: + parts.append(imp) + if std_imports and local_imports: + parts.append("") + for imp in local_imports: parts.append(imp) # Typing imports (grouped by module) diff --git a/src/capnp_stubgen/gen_interface.py b/src/capnp_stubgen/gen_interface.py new file mode 100644 index 0000000..cae3244 --- /dev/null +++ b/src/capnp_stubgen/gen_interface.py @@ -0,0 +1,150 @@ +"""Interface (RPC) type stub generation. + +For each Cap'n Proto interface, generates: +1. Main class — method stubs with typed params/results +2. Client class — RPC client +3. Server class — async method handlers + +Before generation, ``fixup_interface_methods()`` reassigns method +param/result structs to their interface's parent scope so they +appear as nested types. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from .utils import type_id_to_int + +if TYPE_CHECKING: + from .emitter import Emitter + from .models import TypeInfo, TypeRegistry + + +def fixup_interface_methods(registry: TypeRegistry) -> None: + """Reassign method param/result structs to their interface parent. + + Cap'n Proto method param/result structs have ``scopeId=0`` (detached + from the interface scope). This makes them appear as nested types of + the interface by setting their ``parent_type_id`` and adjusting their + ``scoped_name``. + """ + for type_info in list(registry._types.values()): + if type_info.kind != "interface": + continue + + node: dict = type_info.node + for method in node.get("interface", {}).get("methods", []): + mname = method["name"] + param_id = type_id_to_int(method["paramStructType"]) + result_id = type_id_to_int(method["resultStructType"]) + + for sid, suffix in [(param_id, "Params"), (result_id, "Results")]: + sinfo = registry.get(sid) + if sinfo is not None: + sinfo.parent_type_id = type_info.type_id + sinfo.scoped_name = ( + f"{type_info.scoped_name}.{mname}.{suffix}" + ) + # Override name for class nesting + sinfo._display_name = ( + f"{mname}{suffix}" + ) + + +def generate_interface( + emitter: Emitter, + type_info: TypeInfo, + registry: TypeRegistry, +) -> None: + """Generate interface classes for a Cap'n Proto interface.""" + node: dict = type_info.node + iface_body: dict = node.get("interface", {}) + name = type_info.name + methods: list[dict] = iface_body.get("methods", []) + + # Build method info list + method_infos: list[dict] = [] + for method in methods: + param_id = type_id_to_int(method["paramStructType"]) + result_id = type_id_to_int(method["resultStructType"]) + param_info = registry.get(param_id) + result_info = registry.get(result_id) + method_infos.append({ + "name": method["name"], + "param_info": param_info, + "result_info": result_info, + }) + + # ── Main interface class ─────────────────────────────────────────── + emitter.begin_class(name) + + for mi in method_infos: + mname = mi["name"] + param_type = _method_param_type(mi, name) + result_type = _method_result_type(mi, name) + + emitter.add_method( + mname, + params=["self", f"_params: {param_type}"], + return_type=result_type, + ) + + # Method param/result structs nested inside the interface + from .gen_struct import generate_struct as gen_struct_fn + + nested = registry.get_children(type_info.type_id) + struct_children = [nt for nt in nested if nt.kind == "struct"] + for nt in struct_children: + emitter.add_blank_line() + gen_struct_fn(emitter, nt, registry) + + emitter.end_class() + + # ── Client class ─────────────────────────────────────────────────── + client_name = name + "Client" + emitter.begin_class(client_name, bases=[name]) + emitter.add_method( + "_new_client", + params=["self", "server"], + return_type=client_name, + ) + emitter.end_class() + + # ── Server class ─────────────────────────────────────────────────── + server_name = name + "Server" + emitter.begin_class(server_name) + + for mi in method_infos: + mname = mi["name"] + param_type = _method_param_type(mi, name) + + emitter.add_method( + mname, + params=["self", f"_params: {param_type}", "_context"], + return_type=None, + ) + + emitter.end_class() + + +def _method_param_type( + mi: dict, + iface_name: str, +) -> str: + """Get the type string for a method's params.""" + param_info = mi.get("param_info") + if param_info is not None: + return f"{iface_name}.{param_info._display_name or param_info.name}Reader" + return f"{iface_name}.{mi['name']}ParamsReader" + + +def _method_result_type( + mi: dict, + iface_name: str, +) -> str: + """Get the type string for a method's results.""" + result_info = mi.get("result_info") + if result_info is not None: + return f"{iface_name}.{result_info._display_name or result_info.name}Reader" + return f"{iface_name}.{mi['name']}ResultsReader" diff --git a/src/capnp_stubgen/gen_struct.py b/src/capnp_stubgen/gen_struct.py index 33fb578..7600068 100644 --- a/src/capnp_stubgen/gen_struct.py +++ b/src/capnp_stubgen/gen_struct.py @@ -34,7 +34,7 @@ def generate_struct( """Generate the class triple for a Cap'n Proto struct.""" node: dict = type_info.node struct_body: dict = node.get("struct", {}) - name = type_info.name + name = type_info._display_name or type_info.name discriminant_count = struct_body.get("discriminantCount", 0) nested_types = registry.get_children(type_info.type_id) diff --git a/src/capnp_stubgen/models.py b/src/capnp_stubgen/models.py index d69a3bc..a4ed8c3 100644 --- a/src/capnp_stubgen/models.py +++ b/src/capnp_stubgen/models.py @@ -108,6 +108,10 @@ class TypeInfo: file_id: int | None = None """The ID of the file node that contains this type (for import resolution).""" + # May be overridden for method param/result structs + _display_name: str | None = None + """Override for the display name (used for method param/result structs).""" + class TypeRegistry: """Maps Cap'n Proto type IDs to ``TypeInfo`` entries. @@ -120,6 +124,8 @@ class TypeRegistry: self._type_id_to_file_id: dict[int, int] = {} # node ID → file ID mapping for import resolution self._node_to_file: dict[int, int] = {} + # file ID → module name mapping for cross-file imports + self._file_module_names: dict[int, str] = {} def register(self, info: TypeInfo) -> None: """Register a type in the registry.""" @@ -143,6 +149,14 @@ class TypeRegistry: """Get the file ID that contains a given type (for import resolution).""" return self._node_to_file.get(type_id) + def set_file_module_name(self, file_id: int, module_name: str) -> None: + """Record the Python module name for a file node.""" + self._file_module_names[file_id] = module_name + + def get_file_module_name(self, file_id: int) -> str | None: + """Get the Python module name for a file node.""" + return self._file_module_names.get(file_id) + def get_children(self, parent_id: int) -> list[TypeInfo]: """Get all types whose ``scopeId`` matches the given parent ID.""" return [ diff --git a/src/capnp_stubgen/plugin.py b/src/capnp_stubgen/plugin.py index d42adb8..06a4c12 100644 --- a/src/capnp_stubgen/plugin.py +++ b/src/capnp_stubgen/plugin.py @@ -18,6 +18,7 @@ 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 @@ -50,6 +51,22 @@ def _main_impl() -> int: # ── 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: @@ -106,6 +123,9 @@ def _generate_file( """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} @@ -120,7 +140,114 @@ def _generate_file( 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)