Type-level cleanup/evolution in workflow caller

Author: dandavison

temporalio/nexus/_decorators.py

@@ -4,16 +4,17 @@
     Awaitable,
     Callable,
     Optional,
+    TypeVar,
     Union,
     overload,
 )
 
 import nexusrpc
+from nexusrpc import InputT, OutputT
 from nexusrpc.handler import (
     OperationHandler,
     StartOperationContext,
 )
-from nexusrpc.types import InputT, OutputT, ServiceHandlerT
 
 from temporalio.nexus._operation_context import WorkflowRunOperationContext
 from temporalio.nexus._operation_handlers import (
@@ -27,6 +28,8 @@
     get_workflow_run_start_method_input_and_output_type_annotations,
 )
 
+ServiceHandlerT = TypeVar("ServiceHandlerT")
+
 
 @overload
 def workflow_run_operation(

temporalio/nexus/_operation_handlers.py

@@ -8,7 +8,11 @@
     Type,
 )
 
-from nexusrpc import OperationInfo
+from nexusrpc import (
+    InputT,
+    OperationInfo,
+    OutputT,
+)
 from nexusrpc.handler import (
     CancelOperationContext,
     FetchOperationInfoContext,
@@ -19,10 +23,6 @@
     StartOperationContext,
     StartOperationResultAsync,
 )
-from nexusrpc.types import (
-    InputT,
-    OutputT,
-)
 
 from temporalio import client
 from temporalio.nexus._operation_context import (

temporalio/nexus/_token.py

@@ -5,7 +5,7 @@
 from dataclasses import dataclass
 from typing import Any, Generic, Literal, Optional, Type
 
-from nexusrpc.types import OutputT
+from nexusrpc import OutputT
 
 from temporalio import client
 

temporalio/nexus/_util.py

@@ -10,13 +10,13 @@
     Callable,
     Optional,
     Type,
+    TypeVar,
     Union,
 )
 
-from nexusrpc.types import (
+from nexusrpc import (
     InputT,
     OutputT,
-    ServiceHandlerT,
 )
 
 from temporalio.nexus._operation_context import WorkflowRunOperationContext
@@ -25,6 +25,8 @@
     WorkflowHandle as WorkflowHandle,
 )
 
+ServiceHandlerT = TypeVar("ServiceHandlerT")
+
 
 def get_workflow_run_start_method_input_and_output_type_annotations(
     start: Callable[

temporalio/worker/_interceptor.py

@@ -19,10 +19,7 @@
 )
 
 import nexusrpc.handler
-from nexusrpc.types import (
-    InputT,
-    OutputT,
-)
+from nexusrpc import InputT, OutputT
 
 import temporalio.activity
 import temporalio.api.common.v1
@@ -464,7 +461,7 @@ def start_local_activity(
         return self.next.start_local_activity(input)
 
     async def start_nexus_operation(
-        self, input: StartNexusOperationInput
-    ) -> temporalio.workflow.NexusOperationHandle[Any]:
+        self, input: StartNexusOperationInput[InputT, OutputT]
+    ) -> temporalio.workflow.NexusOperationHandle[OutputT]:
         """Called for every :py:func:`temporalio.workflow.start_nexus_operation` call."""
         return await self.next.start_nexus_operation(input)

temporalio/worker/_workflow_instance.py

@@ -44,6 +44,7 @@
 )
 
 import nexusrpc.handler
+from nexusrpc import InputT, OutputT
 from typing_extensions import Self, TypeAlias, TypedDict
 
 import temporalio.activity
@@ -1498,12 +1499,12 @@ async def workflow_start_nexus_operation(
         self,
         endpoint: str,
         service: str,
-        operation: Union[nexusrpc.Operation[I, O], str, Callable[..., Any]],
+        operation: Union[nexusrpc.Operation[InputT, OutputT], str, Callable[..., Any]],
         input: Any,
-        output_type: Optional[Type[O]] = None,
+        output_type: Optional[Type[OutputT]] = None,
         schedule_to_close_timeout: Optional[timedelta] = None,
         headers: Optional[Mapping[str, str]] = None,
-    ) -> temporalio.workflow.NexusOperationHandle[Any]:
+    ) -> temporalio.workflow.NexusOperationHandle[OutputT]:
         # start_nexus_operation
         return await self._outbound.start_nexus_operation(
             StartNexusOperationInput(
@@ -1824,8 +1825,8 @@ async def run_child() -> Any:
                 apply_child_cancel_error()
 
     async def _outbound_start_nexus_operation(
-        self, input: StartNexusOperationInput
-    ) -> _NexusOperationHandle[Any]:
+        self, input: StartNexusOperationInput[Any, OutputT]
+    ) -> _NexusOperationHandle[OutputT]:
         # A Nexus operation handle contains two futures: self._start_fut is resolved as a
         # result of the Nexus operation starting (activation job:
         # resolve_nexus_operation_start), and self._result_fut is resolved as a result of
@@ -1840,9 +1841,9 @@ async def _outbound_start_nexus_operation(
         # and start will be resolved with an operation token). See comments in
         # tests/worker/test_nexus.py for worked examples of the evolution of the resulting
         # handle state machine in the sync and async Nexus response cases.
-        handle: _NexusOperationHandle
+        handle: _NexusOperationHandle[OutputT]
 
-        async def operation_handle_fn() -> Any:
+        async def operation_handle_fn() -> OutputT:
             while True:
                 try:
                     return await asyncio.shield(handle._result_fut)
@@ -2601,8 +2602,8 @@ async def start_child_workflow(
         return await self._instance._outbound_start_child_workflow(input)
 
     async def start_nexus_operation(
-        self, input: StartNexusOperationInput
-    ) -> temporalio.workflow.NexusOperationHandle[Any]:
+        self, input: StartNexusOperationInput[Any, OutputT]
+    ) -> _NexusOperationHandle[OutputT]:
         return await self._instance._outbound_start_nexus_operation(input)
 
     def start_local_activity(
@@ -2991,27 +2992,23 @@ async def cancel(self) -> None:
         await self._instance._cancel_external_workflow(command)
 
 
-I = TypeVar("I")
-O = TypeVar("O")
-
-
 # TODO(dan): are we sure we don't want to inherit from asyncio.Task as ActivityHandle and
 # ChildWorkflowHandle do? I worry that we should provide .done(), .result(), .exception()
 # etc for consistency.
-class _NexusOperationHandle(temporalio.workflow.NexusOperationHandle[O]):
+class _NexusOperationHandle(temporalio.workflow.NexusOperationHandle[OutputT]):
     def __init__(
         self,
         instance: _WorkflowInstanceImpl,
         seq: int,
-        input: StartNexusOperationInput,
-        fn: Coroutine[Any, Any, O],
+        input: StartNexusOperationInput[Any, OutputT],
+        fn: Coroutine[Any, Any, OutputT],
     ):
         self._instance = instance
         self._seq = seq
         self._input = input
         self._task = asyncio.Task(fn)
         self._start_fut: asyncio.Future[Optional[str]] = instance.create_future()
-        self._result_fut: asyncio.Future[Optional[O]] = instance.create_future()
+        self._result_fut: asyncio.Future[Optional[OutputT]] = instance.create_future()
 
     @property
     def operation_token(self) -> Optional[str]:
@@ -3025,10 +3022,10 @@ def operation_token(self) -> Optional[str]:
         except BaseException:
             return None
 
-    async def result(self) -> O:
+    async def result(self) -> OutputT:
         return await self._task
 
-    def __await__(self) -> Generator[Any, Any, O]:
+    def __await__(self) -> Generator[Any, Any, OutputT]:
         return self._task.__await__()
 
     def __repr__(self) -> str:
@@ -3045,7 +3042,7 @@ def _resolve_start_success(self, operation_token: Optional[str]) -> None:
         # We intentionally let this error if already done
         self._start_fut.set_result(operation_token)
 
-    def _resolve_success(self, result: Any) -> None:
+    def _resolve_success(self, result: OutputT) -> None:
         # We intentionally let this error if already done
         self._result_fut.set_result(result)
 

temporalio/workflow.py

@@ -42,6 +42,7 @@
 
 import nexusrpc
 import nexusrpc.handler
+from nexusrpc import InputT, OutputT
 from typing_extensions import (
     Concatenate,
     Literal,
@@ -854,12 +855,12 @@ async def workflow_start_nexus_operation(
         self,
         endpoint: str,
         service: str,
-        operation: Union[nexusrpc.Operation[I, O], str, Callable[..., Any]],
+        operation: Union[nexusrpc.Operation[InputT, OutputT], str, Callable[..., Any]],
         input: Any,
-        output_type: Optional[Type[O]] = None,
+        output_type: Optional[Type[OutputT]] = None,
         schedule_to_close_timeout: Optional[timedelta] = None,
         headers: Optional[Mapping[str, str]] = None,
-    ) -> NexusOperationHandle[Any]: ...
+    ) -> NexusOperationHandle[OutputT]: ...
 
     @abstractmethod
     def workflow_time_ns(self) -> int: ...
@@ -4383,14 +4384,8 @@ async def execute_child_workflow(
     return await handle
 
 
-# TODO(nexus-prerelease): use types from nexusrpc
-I = TypeVar("I")
-O = TypeVar("O")
-S = TypeVar("S")
-
-
 # TODO(nexus-prerelease): ABC / inherit from asyncio.Task?
-class NexusOperationHandle(Generic[O]):
+class NexusOperationHandle(Generic[OutputT]):
     def cancel(self) -> bool:
         # TODO(nexus-prerelease): docstring
         """
@@ -4404,7 +4399,7 @@ def cancel(self) -> bool:
         """
         raise NotImplementedError
 
-    def __await__(self) -> Generator[Any, Any, O]:
+    def __await__(self) -> Generator[Any, Any, OutputT]:
         raise NotImplementedError
 
     # TODO(nexus-prerelease): check SDK-wide consistency for @property vs nullary accessor methods.
@@ -4416,13 +4411,13 @@ def operation_token(self) -> Optional[str]:
 async def start_nexus_operation(
     endpoint: str,
     service: str,
-    operation: Union[nexusrpc.Operation[I, O], str, Callable[..., Any]],
+    operation: Union[nexusrpc.Operation[InputT, OutputT], str, Callable[..., Any]],
     input: Any,
     *,
-    output_type: Optional[Type[O]] = None,
+    output_type: Optional[Type[OutputT]] = None,
     schedule_to_close_timeout: Optional[timedelta] = None,
     headers: Optional[Mapping[str, str]] = None,
-) -> NexusOperationHandle[Any]:
+) -> NexusOperationHandle[OutputT]:
     """Start a Nexus operation and return its handle.
 
     Args:
@@ -5161,17 +5156,13 @@ def _to_proto(self) -> temporalio.bridge.proto.common.VersioningIntent.ValueType
 
 # Nexus
 
+ServiceT = TypeVar("ServiceT")
 
-class NexusClient(Generic[S]):
+
+class NexusClient(Generic[ServiceT]):
     def __init__(
         self,
-        service: Union[
-            # TODO(nexus-prerelease): Type[S] is modeling the interface case as well the impl case, but
-            # the typevar S is used below only in the impl case. I think this is OK, but
-            # think about it again before deleting this TODO.
-            Type[S],
-            str,
-        ],
+        service: Union[Type[ServiceT], str],
         *,
         endpoint: str,
     ) -> None:
@@ -5194,13 +5185,13 @@ def __init__(
     # TODO(nexus-prerelease): should it be an error to use a reference to a method on a class other than that supplied?
     async def start_operation(
         self,
-        operation: Union[nexusrpc.Operation[I, O], str, Callable[..., Any]],
-        input: I,
+        operation: Union[nexusrpc.Operation[InputT, OutputT], str, Callable[..., Any]],
+        input: InputT,
         *,
-        output_type: Optional[Type[O]] = None,
+        output_type: Optional[Type[OutputT]] = None,
         schedule_to_close_timeout: Optional[timedelta] = None,
         headers: Optional[Mapping[str, str]] = None,
-    ) -> NexusOperationHandle[O]:
+    ) -> NexusOperationHandle[OutputT]:
         return await temporalio.workflow.start_nexus_operation(
             endpoint=self._endpoint,
             service=self._service_name,
@@ -5214,14 +5205,14 @@ async def start_operation(
     # TODO(nexus-prerelease): overloads: no-input, ret type
     async def execute_operation(
         self,
-        operation: Union[nexusrpc.Operation[I, O], str, Callable[..., Any]],
-        input: I,
+        operation: Union[nexusrpc.Operation[InputT, OutputT], str, Callable[..., Any]],
+        input: InputT,
         *,
-        output_type: Optional[Type[O]] = None,
+        output_type: Optional[Type[OutputT]] = None,
         schedule_to_close_timeout: Optional[timedelta] = None,
         headers: Optional[Mapping[str, str]] = None,
-    ) -> O:
-        handle: NexusOperationHandle[O] = await self.start_operation(
+    ) -> OutputT:
+        handle: NexusOperationHandle[OutputT] = await self.start_operation(
             operation,
             input,
             output_type=output_type,