Fix typing issues

Author: kirkbrauer

packages/jumpstarter/jumpstarter/client/core.py

@@ -2,6 +2,8 @@
 Base classes for drivers and driver clients
 """
 
+from __future__ import annotations
+
 import logging
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field

packages/jumpstarter/jumpstarter/client/grpc.py

@@ -5,7 +5,6 @@
 from dataclasses import InitVar, dataclass, field
 from datetime import datetime, timedelta
 from types import SimpleNamespace
-from typing import Any
 
 from google.protobuf import duration_pb2, field_mask_pb2, json_format, timestamp_pb2
 from grpc import ChannelConnectivity
@@ -489,7 +488,7 @@ class MultipathExporterStub:
 
     channels: InitVar[list[Channel]]
 
-    __stubs: dict[Channel, Any] = field(init=False, default_factory=OrderedDict)
+    __stubs: dict[Channel, SimpleNamespace] = field(init=False, default_factory=OrderedDict)
 
     def __post_init__(self, channels):
         for channel in channels:

packages/jumpstarter/jumpstarter/common/__init__.py

@@ -1,12 +1,22 @@
 from .enums import ExporterStatus, LogSource
 from .metadata import Metadata
 from .tempfile import TemporarySocket, TemporaryTcpListener, TemporaryUnixListener
+from .types import (
+    AsyncChannel,
+    ControllerStub,
+    ExporterStub,
+    RouterStub,
+)
 
 __all__ = [
+    "AsyncChannel",
+    "ControllerStub",
     "ExporterStatus",
+    "ExporterStub",
     "LogSource",
     "Metadata",
+    "RouterStub",
     "TemporarySocket",
-    "TemporaryUnixListener",
     "TemporaryTcpListener",
+    "TemporaryUnixListener",
 ]

packages/jumpstarter/jumpstarter/common/types.py

@@ -0,0 +1,25 @@
+"""Type aliases for gRPC and Protobuf types."""
+
+from typing import TYPE_CHECKING, TypeAlias
+
+from grpc.aio import Channel
+from jumpstarter_protocol import jumpstarter_pb2_grpc, router_pb2_grpc
+
+# Stub type aliases (the generic Stub classes work for both sync and async)
+ExporterStub: TypeAlias = jumpstarter_pb2_grpc.ExporterServiceStub
+RouterStub: TypeAlias = router_pb2_grpc.RouterServiceStub
+ControllerStub: TypeAlias = jumpstarter_pb2_grpc.ControllerServiceStub
+
+# Channel type alias
+AsyncChannel: TypeAlias = Channel
+
+# Async stub type aliases are only available for type checking (defined in .pyi files)
+if TYPE_CHECKING:
+    pass
+
+__all__ = [
+    "AsyncChannel",
+    "ControllerStub",
+    "ExporterStub",
+    "RouterStub",
+]

packages/jumpstarter/jumpstarter/config/client.py

@@ -120,7 +120,7 @@ class ClientConfigV1Alpha1(BaseSettings):
 
     leases: ClientConfigV1Alpha1Lease = Field(default_factory=ClientConfigV1Alpha1Lease)
 
-    async def channel(self):
+    async def channel(self) -> grpc.aio.Channel:
         if self.endpoint is None or self.token is None:
             raise ConfigurationError("endpoint or token not set in client config")
 

packages/jumpstarter/jumpstarter/config/exporter.py

@@ -212,7 +212,7 @@ async def create_exporter(self):
 
         from jumpstarter.exporter import Exporter
 
-        async def channel_factory():
+        async def channel_factory() -> grpc.aio.Channel:
             if self.endpoint is None or self.token is None:
                 raise ConfigurationError("endpoint or token not set in exporter config")
             credentials = grpc.composite_channel_credentials(