From 756d12f92227c252db2478a39f0163109d24fbc2 Mon Sep 17 00:00:00 2001 From: worldmozara Date: Thu, 9 Jul 2026 15:21:04 +0800 Subject: [PATCH] feat/stubgen: add const and generic support --- src/capnp_stubgen/gen_const.py | 153 +++++++++++++++++++++++++++++ src/capnp_stubgen/gen_generic.py | 72 ++++++++++++++ src/capnp_stubgen/gen_struct.py | 27 ++++- src/capnp_stubgen/plugin.py | 3 + src/capnp_stubgen/type_resolver.py | 30 ++++-- 5 files changed, 272 insertions(+), 13 deletions(-) create mode 100644 src/capnp_stubgen/gen_const.py create mode 100644 src/capnp_stubgen/gen_generic.py diff --git a/src/capnp_stubgen/gen_const.py b/src/capnp_stubgen/gen_const.py new file mode 100644 index 0000000..5fa7bbf --- /dev/null +++ b/src/capnp_stubgen/gen_const.py @@ -0,0 +1,153 @@ +"""Constant type stub generation. + +Handles Cap'n Proto constants declared at file scope or nested within +structs. Primitive values are rendered inline; complex values (struct, +list) use ``...`` as a placeholder. +""" + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING + +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 + + +def generate_const( + emitter: Emitter, + type_info: TypeInfo, + registry: TypeRegistry, +) -> None: + """Generate a typed constant annotation. + + Example outputs:: + + globalInt: int = 12345 + globalText: str = "foobar" + voidConst: None = None + """ + node: dict = type_info.node + const_body: dict = node.get("const", {}) + type_dict: dict = const_body.get("type", {}) + value_dict: dict = const_body.get("value", {}) + + python_type_str = _resolve_const_type(type_dict, registry) + python_value_str = _render_const_value(value_dict, type_dict, registry) + + emitter.add_field( + type_info.name, + python_type_str, + default=python_value_str, + ) + + +# ── type resolution ──────────────────────────────────────────────────────── + + +def _resolve_const_type(type_dict: dict, registry: TypeRegistry) -> str: + """Resolve a constant's type to a Python type annotation string.""" + # Handle enum references + if "enum" in type_dict: + type_id = type_id_to_int(type_dict["enum"]["typeId"]) + info = registry.get(type_id) + if info is not None: + return info.scoped_name + return "Any" + + # Handle struct references (for struct constants we skip the value) + if "struct" in type_dict: + type_id = type_id_to_int(type_dict["struct"]["typeId"]) + info = registry.get(type_id) + if info is not None: + return info.scoped_name + return "Any" + + # Handle list types + if "list" in type_dict: + return "Any" # Simplified for now + + # Primitives + pt = resolve_type(type_dict, registry) + if pt is not None: + return pt.render() + + return "Any" + + +# ── value rendering ───────────────────────────────────────────────────────── + + +def _render_const_value( + value_dict: dict, + type_dict: dict, + registry: TypeRegistry, +) -> str: + """Render a constant's value as a Python literal.""" + + # Void + if "void" in value_dict: + return "None" + + # Bool + if "bool" in value_dict: + return str(value_dict["bool"]) + + # Integers (signed and unsigned) + for ik in ("int8", "int16", "int32", "int64", + "uint8", "uint16", "uint32", "uint64"): + if ik in value_dict: + return str(value_dict[ik]) + + # Floats + if "float32" in value_dict or "float64" in value_dict: + key = "float32" if "float32" in value_dict else "float64" + v = value_dict[key] + if math.isinf(v): + return "float('inf')" if v > 0 else "float('-inf')" + if math.isnan(v): + return "float('nan')" + return repr(v) + + # Text + if "text" in value_dict: + return repr(value_dict["text"]) + + # Data + if "data" in value_dict: + return repr(value_dict["data"]) + + # Enum + if "enum" in value_dict: + ordinal = value_dict["enum"] + return _resolve_enum_value(type_dict, ordinal, registry) + + # Struct / List / AnyPointer — use placeholder + return "..." + + +def _resolve_enum_value( + type_dict: dict, + ordinal: int, + registry: TypeRegistry, +) -> str: + """Resolve an enum ordinal to its enumerant name.""" + if "enum" not in type_dict: + return repr(ordinal) + + type_id = type_id_to_int(type_dict["enum"]["typeId"]) + info = registry.get(type_id) + if info is None: + return repr(ordinal) + + enum_node: dict = info.node + for e in enum_node.get("enum", {}).get("enumerants", []): + if e.get("ordinal") == ordinal or e.get("codeOrder") == ordinal: + return f'"{e["name"]}"' + + # Not found — return ordinal as fallback + return repr(ordinal) diff --git a/src/capnp_stubgen/gen_generic.py b/src/capnp_stubgen/gen_generic.py new file mode 100644 index 0000000..8aa5538 --- /dev/null +++ b/src/capnp_stubgen/gen_generic.py @@ -0,0 +1,72 @@ +"""Generic type support. + +Handles Cap'n Proto generic (parameterized) structs by generating +``TypeVar`` declarations and ``Generic[T]`` base classes in Python stubs. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from .emitter import Emitter + from .models import TypeInfo, TypeRegistry + + +def generate_type_vars( + emitter: Emitter, + type_info: TypeInfo, +) -> list[str]: + """Generate ``TypeVar`` declarations for a generic struct's parameters. + + Returns the list of TypeVar names for use in the ``Generic[...]`` base. + """ + node: dict = type_info.node + params = node.get("parameters", []) + + tv_names: list[str] = [] + for param in params: + tv_name = param["name"] + emitter.add_typing_import("TypeVar") + emitter.add_type_alias(tv_name, f'TypeVar("{tv_name}")') + tv_names.append(tv_name) + + if tv_names: + emitter.add_blank_line() + + return tv_names + + +def resolve_brand( + brand_dict: dict, + registry: TypeRegistry, +) -> list[str] | None: + """Resolve brand bindings to Python type strings. + + Given a brand dict like:: + + {"scopes": [{"bind": [{"type": {"text": null}}], "scopeId": ...}]} + + Returns the resolved Python types, e.g. ``["str"]``. + + Returns ``None`` if the brand cannot be resolved. + """ + from .type_resolver import resolve_type + + scopes = brand_dict.get("scopes", []) + bindings: list[str] = [] + + for scope in scopes: + if "bind" in scope: + for binding in scope["bind"]: + type_dict = binding.get("type", {}) + py_type = resolve_type(type_dict, registry) + if py_type is not None: + bindings.append(py_type.render()) + else: + bindings.append("Any") + elif "inherit" in scope: + # inherit means "use the parent's parameter binding" + bindings.append("Any") # Simplified for now + + return bindings if bindings else None diff --git a/src/capnp_stubgen/gen_struct.py b/src/capnp_stubgen/gen_struct.py index be0735b..33fb578 100644 --- a/src/capnp_stubgen/gen_struct.py +++ b/src/capnp_stubgen/gen_struct.py @@ -12,7 +12,9 @@ from __future__ import annotations from typing import TYPE_CHECKING +from .gen_const import generate_const from .gen_enum import generate_enum +from .gen_generic import generate_type_vars from .type_resolver import resolve_type from .utils import type_id_to_int @@ -68,7 +70,16 @@ def generate_struct( non_union_fields.append((field_name, type_str)) # ── class: main type ─────────────────────────────────────────────── - emitter.begin_class(name) + # Generic type: generate TypeVars and Generic[T] base + tv_names: list[str] = [] + if type_info.generic_params: + tv_names = generate_type_vars(emitter, type_info) + emitter.add_typing_import("Generic") + + if tv_names: + emitter.begin_class(name, bases=[f"Generic[{', '.join(tv_names)}]"]) + else: + emitter.begin_class(name) for field_name, type_str in non_union_fields: emitter.add_field(field_name, type_str) @@ -82,6 +93,12 @@ def generate_struct( emitter.add_blank_line() generate_enum(emitter, nt) + # Nested constants (class-level) + for nt in nested_types: + if nt.kind == "const": + emitter.add_blank_line() + generate_const(emitter, nt, registry) + # Factory methods + to_dict _generate_factory_methods(emitter, name) emitter.add_method("to_dict", params=["self"], return_type="dict") @@ -127,10 +144,10 @@ def _resolve_field_type( 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 + # Use resolve_type to handle brand bindings + py_type = resolve_type(type_dict, registry) + if py_type is not None: + q = py_type.render() return f"{q} | {q}Builder | {q}Reader" return "Any" diff --git a/src/capnp_stubgen/plugin.py b/src/capnp_stubgen/plugin.py index 14364f1..d42adb8 100644 --- a/src/capnp_stubgen/plugin.py +++ b/src/capnp_stubgen/plugin.py @@ -16,6 +16,7 @@ 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_struct import generate_struct from .models import TypeRegistry @@ -119,5 +120,7 @@ def _generate_file( generate_enum(emitter, type_info) elif type_info.kind == "struct": generate_struct(emitter, type_info, registry) + elif type_info.kind == "const": + generate_const(emitter, type_info, registry) return emitter.render() diff --git a/src/capnp_stubgen/type_resolver.py b/src/capnp_stubgen/type_resolver.py index 9ae3332..1b866f5 100644 --- a/src/capnp_stubgen/type_resolver.py +++ b/src/capnp_stubgen/type_resolver.py @@ -63,7 +63,14 @@ def resolve_type( # ── struct ───────────────────────────────────────────────────────── if which == "struct": type_id = type_id_to_int(type_dict["struct"]["typeId"]) - return _resolve_named_type(type_id, registry) + base_type = _resolve_named_type(type_id, registry) + brand = type_dict["struct"].get("brand", {}) + if brand: + from .gen_generic import resolve_brand + bindings = resolve_brand(brand, registry) + if bindings: + base_type.params = [PythonType(name=b) for b in bindings] + return base_type # ── interface ────────────────────────────────────────────────────── if which == "interface": @@ -72,7 +79,7 @@ def resolve_type( # ── anyPointer ───────────────────────────────────────────────────── if which == "anyPointer": - return _resolve_any_pointer(type_dict["anyPointer"]) + return _resolve_any_pointer(type_dict["anyPointer"], registry) return PythonType(name="Any") @@ -121,13 +128,20 @@ def _resolve_named_type(type_id: int, registry: TypeRegistry) -> PythonType: return PythonType(name=info.name) -def _resolve_any_pointer(ap_dict: dict) -> PythonType | None: +def _resolve_any_pointer( + ap_dict: dict, registry: TypeRegistry +) -> PythonType | None: """Resolve an AnyPointer type.""" - sub = _type_which(ap_dict) - if sub == "unconstrained": + if "unconstrained" in ap_dict: return PythonType(name="Any") - if sub == "parameter": - return PythonType(name="Any") # Simplified for Phase 1 - if sub == "implicitMethodParameter": + if "parameter" in ap_dict: + # Look up the generic parameter name from the parent scope + scope_id = type_id_to_int(ap_dict["parameter"]["scopeId"]) + param_index = ap_dict["parameter"]["parameterIndex"] + info = registry.get(scope_id) + if info is not None and param_index < len(info.generic_params): + return PythonType(name=info.generic_params[param_index]) + return PythonType(name="Any") + if "implicitMethodParameter" in ap_dict: return None return PythonType(name="Any")