feat/stubgen: add interface support
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user