feat/stubgen: add struct and enum generator
This commit is contained in:
@@ -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"
|
||||
@@ -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())
|
||||
@@ -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)
|
||||
<regular imports>
|
||||
(blank)
|
||||
<typing imports>
|
||||
(blank)
|
||||
<type definitions>
|
||||
"""
|
||||
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)
|
||||
@@ -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()
|
||||
@@ -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]"])
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
@@ -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")
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user