lint/stubgen: fix typing issues

src/capnp_stubgen/plugin.py

@@ -13,19 +13,16 @@ from __future__ import annotations
 
 import os
 import sys
-from typing import TYPE_CHECKING
+from typing import Any
 
 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 .models import TypeInfo, 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
+from .utils import NodeDict, filename_to_module_name, type_id_to_int
 
 
 def main() -> int:
@@ -38,15 +35,13 @@ def main() -> int:
 
 
 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"]
+    nodes: list[NodeDict] = request_dict["nodes"]
+    requested_files: list[NodeDict] = request_dict["requestedFiles"]
 
     # ── 2. Build type registry ─────────────────────────────────────────
     registry = build_type_registry(nodes, requested_files)
@@ -79,7 +74,7 @@ def _main_impl() -> int:
             output_path = f"{module_name}.pyi"
 
             with open(output_path, "w") as f:
-                f.write(content)
+                _ = f.write(content)
 
         except Exception as exc:
             print(
@@ -100,13 +95,13 @@ def _read_request() -> Any:
     Uses pycapnp's ``SchemaParser`` to dynamically load ``schema.capnp``
     (shipped with pycapnp), then casts the stdin message.
     """
-    import capnp
+    import capnp  # pyright: ignore[reportMissingTypeStubs]
 
     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()
+    parser = capnp.SchemaParser()  # pyright: ignore[reportAttributeAccessIssue]
     schema_module = parser.load(schema_path, imports=[import_base])
 
     return schema_module.CodeGeneratorRequest.read(sys.stdin.buffer)
@@ -130,7 +125,7 @@ def _generate_file(
 
     _ORDER = {"enum": 0, "struct": 1, "const": 2, "interface": 3}
 
-    def _sort_key(info):
+    def _sort_key(info: TypeInfo) -> tuple[int, str]:
         return (_ORDER.get(info.kind, 99), info.name)
 
     top_level.sort(key=_sort_key)
@@ -163,12 +158,12 @@ def _resolve_cross_imports(
 
     # Collect all types belonging to this file
     own_types: set[int] = set()
-    for info in registry._types.values():
+    for info in registry.all_types():
         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():
+    for info in registry.all_types():
         if info.file_id != file_id:
             continue
 
@@ -196,7 +191,7 @@ def _resolve_cross_imports(
         all_imports: list[str] = []
         for tname in sorted(type_names):
             ref_info = None
-            for info in registry._types.values():
+            for info in registry.all_types():
                 if info.scoped_name == tname and info.file_id != file_id:
                     ref_info = info
                     break
@@ -226,7 +221,7 @@ def _collect_type_references(info: TypeInfo, registry: TypeRegistry) -> set[int]
     return refs
 
 
-def _collect_field_references(field: dict, refs: set[int], registry: TypeRegistry) -> None:
+def _collect_field_references(field: NodeDict, 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)
@@ -234,7 +229,7 @@ def _collect_field_references(field: dict, refs: set[int], registry: TypeRegistr
         refs.add(type_id_to_int(field["group"]["typeId"]))
 
 
-def _collect_type_references_from_dict(type_dict: dict, refs: set[int]) -> None:
+def _collect_type_references_from_dict(type_dict: NodeDict, 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"]))