From b44b420a8857493be4f3ef1f8a94b5d9f7a2b447 Mon Sep 17 00:00:00 2001 From: worldmozara Date: Thu, 9 Jul 2026 15:06:39 +0800 Subject: [PATCH] feat/stubgen: add struct and enum generator --- src/capnp_stubgen/__init__.py | 7 + src/capnp_stubgen/__main__.py | 12 ++ src/capnp_stubgen/emitter.py | 197 +++++++++++++++++++ src/capnp_stubgen/gen_enum.py | 35 ++++ src/capnp_stubgen/gen_struct.py | 303 +++++++++++++++++++++++++++++ src/capnp_stubgen/models.py | 162 +++++++++++++++ src/capnp_stubgen/plugin.py | 123 ++++++++++++ src/capnp_stubgen/schema_walker.py | 134 +++++++++++++ src/capnp_stubgen/type_resolver.py | 133 +++++++++++++ src/capnp_stubgen/utils.py | 41 ++++ 10 files changed, 1147 insertions(+) create mode 100644 src/capnp_stubgen/__init__.py create mode 100644 src/capnp_stubgen/__main__.py create mode 100644 src/capnp_stubgen/emitter.py create mode 100644 src/capnp_stubgen/gen_enum.py create mode 100644 src/capnp_stubgen/gen_struct.py create mode 100644 src/capnp_stubgen/models.py create mode 100644 src/capnp_stubgen/plugin.py create mode 100644 src/capnp_stubgen/schema_walker.py create mode 100644 src/capnp_stubgen/type_resolver.py create mode 100644 src/capnp_stubgen/utils.py diff --git a/src/capnp_stubgen/__init__.py b/src/capnp_stubgen/__init__.py new file mode 100644 index 0000000..865c85c --- /dev/null +++ b/src/capnp_stubgen/__init__.py @@ -0,0 +1,7 @@ +"""Cap'n Proto Python type stub generator (capnp compile plugin). + +Generates .pyi type stub files for Cap'n Proto schemas, providing IDE +autocompletion and static type checking support via Pyright/Pylance. +""" + +__version__ = "0.1.0" diff --git a/src/capnp_stubgen/__main__.py b/src/capnp_stubgen/__main__.py new file mode 100644 index 0000000..9458d0c --- /dev/null +++ b/src/capnp_stubgen/__main__.py @@ -0,0 +1,12 @@ +"""Allow running as ``python -m capnp_stubgen``. + +This is mainly useful for running the plugin manually (without capnp compile) +if you have a pre-recorded CodeGeneratorRequest. +""" + +from __future__ import annotations + +from .plugin import main + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/capnp_stubgen/emitter.py b/src/capnp_stubgen/emitter.py new file mode 100644 index 0000000..bf9f02e --- /dev/null +++ b/src/capnp_stubgen/emitter.py @@ -0,0 +1,197 @@ +""".pyi type stub file builder. + +The ``Emitter`` class constructs the content of a ``.pyi`` type stub file: +managing imports, indentation, class/method/field declarations, and type aliases. + +Imports are collected during code generation and rendered before type +definitions in the final output. +""" + +from __future__ import annotations + +_INDENT = " " # 4 spaces + + +class Emitter: + """Builds the string content for a single ``.pyi`` stub file. + + Usage:: + + emitter = Emitter("myschema.capnp") + emitter.begin_class("Person") + emitter.add_field("name", "str") + emitter.end_class() + content = emitter.render() + """ + + def __init__(self, source_filename: str = "") -> None: + self._source_filename = source_filename + self._header_lines: list[str] = [] + self._body_lines: list[str] = [] + self._indent_level: int = 0 + self._imports: set[str] = set() + self._typing_imports: dict[str, set[str]] = {} # module → {names} + + self._build_header() + + # ── header ────────────────────────────────────────────────────────────── + + def _build_header(self) -> None: + """Pre-build the header lines (docstring + __future__).""" + if self._source_filename: + self._header_lines.append( + f'"""Auto-generated type stub for {self._source_filename}."""' + ) + else: + self._header_lines.append('"""Auto-generated type stub."""') + self._header_lines.append("") + self._header_lines.append("from __future__ import annotations") + + # ── public API ───────────────────────────────────────────────────────── + + def begin_class(self, name: str, bases: list[str] | None = None) -> None: + """Open a class definition block.""" + if bases: + self._push(f"class {name}({', '.join(bases)}):") + else: + self._push(f"class {name}:") + self._indent_level += 1 + + def end_class(self) -> None: + """Close a class definition block. Emits two blank lines for spacing.""" + self._indent_level = max(0, self._indent_level - 1) + self._push("") + self._push("") + + def add_field( + self, name: str, type_str: str, default: str | None = None + ) -> None: + """Emit a field annotation.""" + if default: + self._push(f"{name}: {type_str} = {default}") + else: + self._push(f"{name}: {type_str}") + + def add_method( + self, + name: str, + params: list[str] | None = None, + return_type: str | None = None, + *, + decorators: list[str] | None = None, + body: str = "...", + ) -> None: + """Emit a method stub.""" + if decorators: + for dec in decorators: + self._push(f"@{dec}") + + params_str = ", ".join(params) if params else "" + if return_type: + sig = f"def {name}({params_str}) -> {return_type}: {body}" + else: + sig = f"def {name}({params_str}): {body}" + self._push(sig) + + def add_static_method( + self, + name: str, + params: list[str] | None = None, + return_type: str | None = None, + *, + decorators: list[str] | None = None, + body: str = "...", + ) -> None: + """Emit a @staticmethod method stub.""" + self._push("@staticmethod") + self.add_method( + name, params, return_type, + decorators=decorators, + body=body, + ) + + def add_decorator(self, decorator: str) -> None: + """Emit a decorator line (for use with @overload, etc.).""" + self._push(f"@{decorator}") + + def add_type_alias(self, name: str, value: str) -> None: + """Emit a top-level type alias like ``Gender = Literal["male", "female"]``.""" + self._push(f"{name} = {value}") + + def add_blank_line(self) -> None: + """Emit an empty line (useful for spacing between definitions).""" + self._push("") + + # ── import management ────────────────────────────────────────────────── + + def add_import(self, line: str) -> None: + """Add a full import line (e.g. ``"import os"``, ``"from foo import bar"``).""" + self._imports.add(line) + + def add_typing_import(self, name: str, module: str = "typing") -> None: + """Add a symbol to import from a module (default: ``typing``). + + Multiple calls for the same module are merged:: + + e.add_typing_import("Sequence") + e.add_typing_import("Literal") + # → ``from typing import Literal, Sequence`` + """ + if module not in self._typing_imports: + self._typing_imports[module] = set() + self._typing_imports[module].add(name) + + # ── rendering ────────────────────────────────────────────────────────── + + def render(self) -> str: + """Return the complete .pyi file content as a string. + + The output is structured as:: + + docstring + __future__ import + (blank) + + (blank) + + (blank) + + """ + parts: list[str] = list(self._header_lines) + + # Blank line between header and imports + parts.append("") + + # Regular imports + if self._imports: + for imp in sorted(self._imports): + parts.append(imp) + + # Typing imports (grouped by module) + if self._typing_imports: + if self._imports: + parts.append("") + # Sort typing modules for determinism + for module in sorted(self._typing_imports): + names = self._typing_imports[module] + if names: + joined = ", ".join(sorted(names)) + parts.append(f"from {module} import {joined}") + + # Blank line between imports and body + parts.append("") + parts.append("") + + # Body (type definitions) + parts.extend(self._body_lines) + + return "\n".join(parts) + "\n" + + # ── internal ─────────────────────────────────────────────────────────── + + def _push(self, line: str) -> None: + """Append a line to the body at the current indentation level.""" + if line == "": + self._body_lines.append("") + else: + self._body_lines.append(_INDENT * self._indent_level + line) diff --git a/src/capnp_stubgen/gen_enum.py b/src/capnp_stubgen/gen_enum.py new file mode 100644 index 0000000..f322436 --- /dev/null +++ b/src/capnp_stubgen/gen_enum.py @@ -0,0 +1,35 @@ +"""Enum type stub generation. + +Cap'n Proto enums are represented as ``Literal[...]`` type aliases. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from .emitter import Emitter + from .models import TypeInfo + + +def generate_enum( + emitter: Emitter, + type_info: TypeInfo, +) -> None: + """Generate a ``Literal[...]`` type alias for a Cap'n Proto enum. + + Example output:: + + Gender = Literal["male", "female", "other"] + """ + node: dict = type_info.node + enum_body: dict = node.get("enum", {}) + enumerants: list[dict] = enum_body.get("enumerants", []) + + literals: list[str] = [] + for e in enumerants: + literals.append(f'"{e["name"]}"') + + emitter.add_typing_import("Literal") + emitter.add_type_alias(type_info.name, f"Literal[{', '.join(literals)}]") + emitter.add_blank_line() diff --git a/src/capnp_stubgen/gen_struct.py b/src/capnp_stubgen/gen_struct.py new file mode 100644 index 0000000..be0735b --- /dev/null +++ b/src/capnp_stubgen/gen_struct.py @@ -0,0 +1,303 @@ +"""Struct type stub generation. + +For each Cap'n Proto struct, generates three Python classes: +1. Main class — field annotations + factory/static methods +2. Reader class — read-only view +3. Builder class — read-write view + +Works with dict nodes from ``CodeGeneratorRequest.to_dict()``. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from .gen_enum import generate_enum +from .type_resolver import resolve_type +from .utils import type_id_to_int + +if TYPE_CHECKING: + from .emitter import Emitter + from .models import TypeInfo, TypeRegistry + + +_NO_DISCRIMINANT = 65535 # 0xffff + + +def generate_struct( + emitter: Emitter, + type_info: TypeInfo, + registry: TypeRegistry, +) -> None: + """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 + + discriminant_count = struct_body.get("discriminantCount", 0) + nested_types = registry.get_children(type_info.type_id) + + non_union_fields: list[tuple[str, str]] = [] + union_fields: list[tuple[str, str]] = [] + + # ── collect fields ───────────────────────────────────────────────── + for field in struct_body.get("fields", []): + field_name: str = field["name"] + disc_value: int = field.get("discriminantValue", _NO_DISCRIMINANT) + + if "slot" in field: + slot = field["slot"] + type_dict = slot.get("type", {}) + field_type_str = _resolve_field_type(type_dict, registry, emitter) + if field_type_str is None: + continue + + if disc_value != _NO_DISCRIMINANT: + union_fields.append((field_name, field_type_str)) + else: + non_union_fields.append((field_name, field_type_str)) + + elif "group" in field: + group_type_id = type_id_to_int(field["group"]["typeId"]) + group_info = registry.get(group_type_id) + if group_info is not None: + type_str = group_info.scoped_name + if disc_value != _NO_DISCRIMINANT: + union_fields.append((field_name, type_str)) + else: + non_union_fields.append((field_name, type_str)) + + # ── class: main type ─────────────────────────────────────────────── + emitter.begin_class(name) + + for field_name, type_str in non_union_fields: + emitter.add_field(field_name, type_str) + + if discriminant_count > 0 and union_fields: + _generate_union_methods(emitter, union_fields) + + # Nested enums first + for nt in nested_types: + if nt.kind == "enum": + emitter.add_blank_line() + generate_enum(emitter, nt) + + # Factory methods + to_dict + _generate_factory_methods(emitter, name) + emitter.add_method("to_dict", params=["self"], return_type="dict") + + # Nested structs + for nt in nested_types: + if nt.kind == "struct": + emitter.add_blank_line() + generate_struct(emitter, nt, registry) + + emitter.end_class() + + # ── class: Reader ────────────────────────────────────────────────── + reader_name = name + "Reader" + emitter.begin_class(reader_name, bases=[name]) + + for field_name, type_str in non_union_fields: + emitter.add_field(field_name, _to_reader_type(type_str)) + + emitter.add_method( + "as_builder", params=["self"], return_type=name + "Builder", + ) + emitter.end_class() + + # ── class: Builder ───────────────────────────────────────────────── + builder_name = name + "Builder" + emitter.begin_class(builder_name, bases=[name]) + + for field_name, type_str in non_union_fields: + emitter.add_field(field_name, _to_builder_type(type_str)) + + _generate_builder_methods(emitter, name) + emitter.end_class() + + +# ── field type resolution ────────────────────────────────────────────────── + + +def _resolve_field_type( + type_dict: dict, registry: TypeRegistry, emitter: Emitter, +) -> str | None: + """Resolve a field's type to a Python type annotation string.""" + which = _type_which(type_dict) + + if which == "struct": + type_id = type_id_to_int(type_dict["struct"]["typeId"]) + info = registry.get(type_id) + if info is not None: + q = info.scoped_name + return f"{q} | {q}Builder | {q}Reader" + return "Any" + + if which == "list": + return _resolve_list_field_type(type_dict["list"], registry, emitter) + + py_type = resolve_type(type_dict, registry) + return py_type.render() if py_type else "Any" + + +def _resolve_list_field_type( + list_dict: dict, registry: TypeRegistry, emitter: Emitter, +) -> str: + """Resolve a ``List(T)`` field type.""" + emitter.add_typing_import("Sequence", "collections.abc") + + depth = 1 + inner = list_dict.get("elementType", {}) + while _type_which(inner) == "list": + depth += 1 + inner = inner.get("list", {}).get("elementType", {}) + + inner_which = _type_which(inner) + + if inner_which == "struct": + type_id = type_id_to_int(inner["struct"]["typeId"]) + info = registry.get(type_id) + if info is not None: + q = info.scoped_name + el_type = f"{q} | {q}Builder | {q}Reader" + else: + el_type = "Any" + elif inner_which == "enum": + type_id = type_id_to_int(inner["enum"]["typeId"]) + info = registry.get(type_id) + el_type = info.scoped_name if info else "Any" + else: + pt = resolve_type(inner, registry) + el_type = pt.render() if pt else "Any" + + result = f"Sequence[{el_type}]" + for _ in range(depth - 1): + result = f"Sequence[{result}]" + + return result + + +def _type_which(type_dict: dict) -> str | None: + """Return the union variant key for a type dict.""" + for key in ("void", "bool", "int8", "int16", "int32", "int64", + "uint8", "uint16", "uint32", "uint64", + "float32", "float64", "text", "data", + "list", "enum", "struct", "interface", "anyPointer"): + if key in type_dict: + return key + return None + + +# ── Reader/Builder type transformation ───────────────────────────────────── + + +def _to_reader_type(type_str: str) -> str: + return _transform_variant(type_str, "Reader") + + +def _to_builder_type(type_str: str) -> str: + return _transform_variant(type_str, "Builder") + + +def _transform_variant(type_str: str, variant: str) -> str: + if " | " not in type_str: + return type_str + + if type_str.startswith("Sequence["): + prefix = "Sequence[" + inner = type_str[len(prefix):-1] + return f"Sequence[{_transform_variant(inner, variant)}]" + + parts = [p.strip() for p in type_str.split("|")] + for p in parts: + if p.endswith(variant): + return p + return parts[0] + + +# ── method generation ────────────────────────────────────────────────────── + + +def _generate_union_methods( + emitter: Emitter, union_fields: list[tuple[str, str]], +) -> None: + emitter.add_typing_import("Literal") + emitter.add_typing_import("overload") + + literal_values = ", ".join(f'"{name}"' for name, _ in union_fields) + emitter.add_method( + "which", params=["self"], + return_type=f"Literal[{literal_values}]", + ) + + for field_name, type_str in union_fields: + emitter.add_decorator("overload") + emitter.add_method( + "init", + params=["self", f'name: Literal["{field_name}"]'], + return_type=type_str, + ) + + +def _generate_factory_methods(emitter: Emitter, name: str) -> None: + emitter.add_typing_import("Iterator", "collections.abc") + emitter.add_typing_import("contextmanager", "contextlib") + + reader_name = name + "Reader" + builder_name = name + "Builder" + + emitter.add_blank_line() + emitter.add_static_method( + "new_message", params=["**kwargs"], return_type=builder_name, + ) + emitter.add_blank_line() + + emitter.add_decorator("staticmethod") + emitter.add_decorator("contextmanager") + emitter.add_method( + "from_bytes", + params=[ + "data: bytes", + "traversal_limit_in_words: int | None = ...", + "nesting_limit: int | None = ...", + ], + return_type=f"Iterator[{reader_name}]", + ) + emitter.add_blank_line() + + emitter.add_static_method( + "from_bytes_packed", + params=[ + "data: bytes", + "traversal_limit_in_words: int | None = ...", + "nesting_limit: int | None = ...", + ], + return_type=reader_name, + ) + + +def _generate_builder_methods(emitter: Emitter, name: str) -> None: + reader_name = name + "Reader" + builder_name = name + "Builder" + + emitter.add_blank_line() + emitter.add_static_method( + "from_dict", params=["dictionary: dict"], return_type=builder_name, + ) + emitter.add_blank_line() + emitter.add_method("copy", params=["self"], return_type=builder_name) + emitter.add_blank_line() + emitter.add_method("to_bytes", params=["self"], return_type="bytes") + emitter.add_blank_line() + emitter.add_method("to_bytes_packed", params=["self"], return_type="bytes") + emitter.add_blank_line() + emitter.add_method("to_segments", params=["self"], return_type="list[bytes]") + emitter.add_blank_line() + emitter.add_method("as_reader", params=["self"], return_type=reader_name) + emitter.add_blank_line() + + emitter.add_typing_import("IO", "typing") + emitter.add_static_method("write", params=["file: IO[bytes]"]) + emitter.add_blank_line() + emitter.add_static_method("write_packed", params=["file: IO[bytes]"]) diff --git a/src/capnp_stubgen/models.py b/src/capnp_stubgen/models.py new file mode 100644 index 0000000..d69a3bc --- /dev/null +++ b/src/capnp_stubgen/models.py @@ -0,0 +1,162 @@ +"""Data models for capnp-stubgen. + +Defines the core data structures used throughout the stub generator: +``PythonType`` for resolved Python type annotations, +``TypeInfo`` for per-type metadata, +and ``TypeRegistry`` for the global type ID → TypeInfo mapping. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + + +@dataclass +class PythonType: + """A resolved Python type annotation. + + Examples:: + + PythonType(name="int") → ``int`` + PythonType(name="Sequence", params=[...]) → ``Sequence[Person]`` + PythonType(name="Person", qualifiers=["Outer"]) → ``Outer.Person`` + """ + + name: str + """Base type name, e.g. ``"int"``, ``"Person"``, ``"Sequence"``.""" + + params: list[PythonType] = field(default_factory=list) + """Generic type parameters, e.g. for ``Sequence[Person]``, params=[Person].""" + + qualifiers: list[str] = field(default_factory=list) + """Scope qualifiers for nested types, e.g. ``["Outer"]`` → Outer.Person.""" + + def render(self) -> str: + """Render this type as a Python type annotation string.""" + if self.name == "None": + return "None" + if self.name == "...": + return "..." + + base = self.name + if self.qualifiers: + base = ".".join(self.qualifiers + [self.name]) + + if self.params: + rendered_params = ", ".join(p.render() for p in self.params) + return f"{base}[{rendered_params}]" + + return base + + def with_reader_suffix(self) -> PythonType: + """Return a new PythonType with 'Reader' appended to the name.""" + return PythonType( + name=self.name + "Reader", + params=list(self.params), + qualifiers=list(self.qualifiers), + ) + + def with_builder_suffix(self) -> PythonType: + """Return a new PythonType with 'Builder' appended to the name.""" + return PythonType( + name=self.name + "Builder", + params=list(self.params), + qualifiers=list(self.qualifiers), + ) + + def union_type(self, *others: PythonType) -> PythonType: + """Create a union type like ``Foo | FooBuilder | FooReader``. + + Returns ``self`` unchanged if no other types are provided. + """ + # We represent unions as a PythonType with params joined by " | " + all_types = [self] + list(others) + rendered = " | ".join(t.render() for t in all_types) + # Wrap in a special marker + return PythonType(name=rendered) + + +@dataclass +class TypeInfo: + """Metadata about a registered Cap'n Proto type.""" + + type_id: int + """The globally unique 64-bit type ID.""" + + name: str + """Short name, e.g. ``"Person"``.""" + + scoped_name: str + """Fully qualified name, e.g. ``"MyFile.Person"``.""" + + node: Any + """The raw ``schema::Node::Reader`` from the CodeGeneratorRequest.""" + + schema: Any + """The Schema object from ``capnp.SchemaLoader`` (may be None for some types).""" + + kind: str + """One of ``"struct"``, ``"enum"``, ``"interface"``, ``"const"``, ``"annotation"``, ``"file"``.""" + + generic_params: list[str] = field(default_factory=list) + """Names of generic type parameters (empty if not generic).""" + + parent_type_id: int | None = None + """The ``scopeId`` of the parent node, or None for top-level types.""" + + file_id: int | None = None + """The ID of the file node that contains this type (for import resolution).""" + + +class TypeRegistry: + """Maps Cap'n Proto type IDs to ``TypeInfo`` entries. + + Provides cross-file type resolution and scope-chain lookups. + """ + + def __init__(self) -> None: + self._types: dict[int, TypeInfo] = {} + self._type_id_to_file_id: dict[int, int] = {} + # node ID → file ID mapping for import resolution + self._node_to_file: dict[int, int] = {} + + def register(self, info: TypeInfo) -> None: + """Register a type in the registry.""" + self._types[info.type_id] = info + + def get(self, type_id: int) -> TypeInfo | None: + """Look up a type by its ID. Returns None if not found.""" + return self._types.get(type_id) + + def get_or_raise(self, type_id: int) -> TypeInfo: + """Look up a type by its ID. Raises KeyError if not found.""" + if type_id not in self._types: + raise KeyError(f"Type ID {type_id:#018x} not found in registry") + return self._types[type_id] + + def set_node_file(self, node_id: int, file_id: int) -> None: + """Record which file a node belongs to (for import resolution).""" + self._node_to_file[node_id] = file_id + + def get_file_id(self, type_id: int) -> int | None: + """Get the file ID that contains a given type (for import resolution).""" + return self._node_to_file.get(type_id) + + def get_children(self, parent_id: int) -> list[TypeInfo]: + """Get all types whose ``scopeId`` matches the given parent ID.""" + return [ + info + for info in self._types.values() + if info.parent_type_id == parent_id + ] + + def get_top_level_types(self, file_id: int) -> list[TypeInfo]: + """Get all top-level types in a file (those whose parent is the file node).""" + return self.get_children(file_id) + + def __len__(self) -> int: + return len(self._types) + + def __contains__(self, type_id: int) -> bool: + return type_id in self._types diff --git a/src/capnp_stubgen/plugin.py b/src/capnp_stubgen/plugin.py new file mode 100644 index 0000000..14364f1 --- /dev/null +++ b/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() diff --git a/src/capnp_stubgen/schema_walker.py b/src/capnp_stubgen/schema_walker.py new file mode 100644 index 0000000..1dcc03a --- /dev/null +++ b/src/capnp_stubgen/schema_walker.py @@ -0,0 +1,134 @@ +"""Schema node traversal and TypeRegistry construction. + +Walks all node dicts from a ``CodeGeneratorRequest.to_dict()``, +building a ``TypeRegistry`` that maps type IDs to ``TypeInfo``. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from .models import TypeInfo, TypeRegistry +from .utils import get_display_name, get_node_kind, type_id_to_int + +if TYPE_CHECKING: + from typing import Any + + +_GENERATABLE_KINDS = frozenset({"struct", "enum", "interface", "const"}) + + +def build_type_registry( + nodes: list[dict], + requested_files: list[dict], +) -> TypeRegistry: + """Build a ``TypeRegistry`` from the to-dict CodeGeneratorRequest nodes.""" + registry = TypeRegistry() + + # Build ID → node lookup + nodes_by_id: dict[int, dict] = {} + for node in nodes: + nodes_by_id[type_id_to_int(node["id"])] = node + + # Associate nodes with files + _build_file_associations(nodes_by_id, requested_files, registry) + + # Register TypeInfo for each generatable kind + for node in nodes: + kind = get_node_kind(node) + if kind not in _GENERATABLE_KINDS: + continue + + type_id = type_id_to_int(node["id"]) + name = get_display_name(node) + scope_id = node.get("scopeId", "0") + parent_id = type_id_to_int(scope_id) if scope_id != "0" else None + + scoped_name = _build_scoped_name(type_id, nodes_by_id) + + generic_params: list[str] = [] + if node.get("isGeneric"): + for param in node.get("parameters", []): + generic_params.append(param.get("name", "")) + + info = TypeInfo( + type_id=type_id, + name=name, + scoped_name=scoped_name, + node=node, + schema=None, + kind=kind, + generic_params=generic_params, + parent_type_id=parent_id, + file_id=_find_containing_file(type_id, nodes_by_id), + ) + registry.register(info) + + return registry + + +# ── file association ─────────────────────────────────────────────────────── + + +def _build_file_associations( + nodes_by_id: dict[int, dict], + requested_files: list[dict], + registry: TypeRegistry, +) -> None: + for node_id in nodes_by_id: + file_id = _find_containing_file(node_id, nodes_by_id) + if file_id is not None: + registry.set_node_file(node_id, file_id) + + +def _find_containing_file( + node_id: int, nodes_by_id: dict[int, dict] +) -> int | None: + visited: set[int] = set() + current_id = node_id + + while current_id not in visited: + visited.add(current_id) + node = nodes_by_id.get(current_id) + if node is None: + break + + if get_node_kind(node) == "file": + return current_id + + scope_id_str = node.get("scopeId", "0") + scope_id = int(scope_id_str) if scope_id_str != "0" else 0 + if scope_id == 0: + break + current_id = scope_id + + return None + + +# ── scoped name ──────────────────────────────────────────────────────────── + + +def _build_scoped_name( + type_id: int, nodes_by_id: dict[int, dict] +) -> str: + parts: list[str] = [] + visited: set[int] = set() + current_id = type_id + + while current_id not in visited: + visited.add(current_id) + node = nodes_by_id.get(current_id) + if node is None: + break + + if get_node_kind(node) != "file": + parts.append(get_display_name(node)) + + scope_id_str = node.get("scopeId", "0") + scope_id = int(scope_id_str) if scope_id_str != "0" else 0 + if scope_id == 0: + break + current_id = scope_id + + parts.reverse() + return ".".join(parts) diff --git a/src/capnp_stubgen/type_resolver.py b/src/capnp_stubgen/type_resolver.py new file mode 100644 index 0000000..9ae3332 --- /dev/null +++ b/src/capnp_stubgen/type_resolver.py @@ -0,0 +1,133 @@ +"""Cap'n Proto type → Python type annotation resolution. + +Maps type dicts (from ``CodeGeneratorRequest.to_dict()``) to ``PythonType``. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from .models import PythonType, TypeRegistry +from .utils import type_id_to_int + +if TYPE_CHECKING: + from typing import Any + + +# Cap'n Proto type key → Python type name +CAPNP_TO_PYTHON: dict[str, str] = { + "void": "None", + "bool": "bool", + "int8": "int", + "int16": "int", + "int32": "int", + "int64": "int", + "uint8": "int", + "uint16": "int", + "uint32": "int", + "uint64": "int", + "float32": "float", + "float64": "float", + "text": "str", + "data": "bytes", +} + + +def resolve_type( + type_dict: dict, + registry: TypeRegistry, +) -> PythonType | None: + """Resolve a type dict to a ``PythonType``. + + The type dict uses union-variant keys, e.g. ``{"int32": {}}``, + ``{"struct": {"typeId": "123"}}``, ``{"list": {"elementType": {...}}}``. + """ + # Find the union variant key + which = _type_which(type_dict) + if which is None: + return PythonType(name="Any") + + # ── primitives ──────────────────────────────────────────────────── + if which in CAPNP_TO_PYTHON: + return PythonType(name=CAPNP_TO_PYTHON[which]) + + # ── list ─────────────────────────────────────────────────────────── + if which == "list": + return _resolve_list_type(type_dict["list"], registry) + + # ── enum ─────────────────────────────────────────────────────────── + if which == "enum": + type_id = type_id_to_int(type_dict["enum"]["typeId"]) + return _resolve_named_type(type_id, registry) + + # ── struct ───────────────────────────────────────────────────────── + if which == "struct": + type_id = type_id_to_int(type_dict["struct"]["typeId"]) + return _resolve_named_type(type_id, registry) + + # ── interface ────────────────────────────────────────────────────── + if which == "interface": + type_id = type_id_to_int(type_dict["interface"]["typeId"]) + return _resolve_named_type(type_id, registry) + + # ── anyPointer ───────────────────────────────────────────────────── + if which == "anyPointer": + return _resolve_any_pointer(type_dict["anyPointer"]) + + return PythonType(name="Any") + + +def _type_which(type_dict: dict) -> str | None: + """Return the union variant key for a type dict.""" + for key in CAPNP_TO_PYTHON: + if key in type_dict: + return key + for key in ("list", "enum", "struct", "interface", "anyPointer"): + if key in type_dict: + return key + return None + + +def _resolve_list_type(list_dict: dict, registry: TypeRegistry) -> PythonType: + """Resolve a List(T) type.""" + depth = 1 + inner = list_dict.get("elementType", {}) + + # Count nesting depth + while _type_which(inner) == "list": + depth += 1 + inner = inner.get("list", {}).get("elementType", {}) + + element_type = resolve_type(inner, registry) + if element_type is None: + element_type = PythonType(name="Any") + + result = PythonType(name="Sequence", params=[element_type]) + for _ in range(depth - 1): + result = PythonType(name="Sequence", params=[result]) + + return result + + +def _resolve_named_type(type_id: int, registry: TypeRegistry) -> PythonType: + """Resolve a struct/enum/interface reference by type ID.""" + info = registry.get(type_id) + if info is None: + return PythonType(name="Any") + + parts = info.scoped_name.split(".") + if len(parts) > 1: + return PythonType(name=parts[-1], qualifiers=parts[:-1]) + return PythonType(name=info.name) + + +def _resolve_any_pointer(ap_dict: dict) -> PythonType | None: + """Resolve an AnyPointer type.""" + sub = _type_which(ap_dict) + if sub == "unconstrained": + return PythonType(name="Any") + if sub == "parameter": + return PythonType(name="Any") # Simplified for Phase 1 + if sub == "implicitMethodParameter": + return None + return PythonType(name="Any") diff --git a/src/capnp_stubgen/utils.py b/src/capnp_stubgen/utils.py new file mode 100644 index 0000000..013ec9e --- /dev/null +++ b/src/capnp_stubgen/utils.py @@ -0,0 +1,41 @@ +"""Utility functions for capnp-stubgen.""" + +import os + + +def filename_to_module_name(filename: str) -> str: + """Convert a .capnp filename to the corresponding Python module name. + + Example: ``path/to/my_file.capnp`` → ``my_file_capnp`` + """ + base = os.path.basename(filename) + if base.endswith(".capnp"): + base = base[:-6] + return base + "_capnp" + + +def get_display_name(node: dict) -> str: + """Extract the short display name from a node dict. + + The node's ``displayName`` field contains a path like + ``path/to/file.capnp:TypeName``. This removes the file prefix. + """ + name: str = node["displayName"] + prefix_len: int = node["displayNamePrefixLength"] + return name[prefix_len:] + + +def get_node_kind(node: dict) -> str: + """Determine the kind of a node from its dict representation. + + Returns one of: "file", "struct", "enum", "interface", "const", "annotation". + """ + for kind in ("file", "struct", "enum", "interface", "const", "annotation"): + if kind in node: + return kind + return "unknown" + + +def type_id_to_int(type_id: str) -> int: + """Convert a type ID string (from dict) to an integer.""" + return int(type_id)