Support client deriving types from operation impl or interface

Author: dandavison

temporalio/nexus/handler.py

@@ -20,14 +20,14 @@
     TypeVar,
 )
 
-import nexusrpc.handler
 from hyperlinked import hyperlinked, print
-from nexusrpc.handler import _ServiceImpl
 from typing_extensions import Concatenate, overload
 
+import nexusrpc.handler
 import temporalio.api.common.v1
 import temporalio.api.enums.v1
 import temporalio.common
+from nexusrpc.handler import _ServiceImpl
 from temporalio.client import (
     Client,
     WorkflowHandle,
@@ -298,7 +298,7 @@ def workflow_operation(
     def factory(service: S) -> WorkflowOperation[I, O]:
         return WorkflowOperation(service, start_method)
 
-    factory.__nexus_operation__ = nexusrpc.handler._NexusOperationDefinition(
+    factory.__nexus_operation__ = nexusrpc.handler.NexusOperationDefinition(
         name=start_method.__name__
     )
     return factory

temporalio/worker/_interceptor.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import concurrent.futures
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from datetime import timedelta
 from typing import (
     Any,
@@ -22,6 +22,8 @@
 )
 
 import nexusrpc
+import nexusrpc.handler
+import nexusrpc.interface
 
 import temporalio.activity
 import temporalio.api.common.v1
@@ -289,6 +291,7 @@ class StartChildWorkflowInput:
     ret_type: Optional[Type]
 
 
+# TODO(dan): Put these in a better location. Type variance?
 I = TypeVar("I")
 O = TypeVar("O")
 
@@ -299,34 +302,53 @@ class StartNexusOperationInput(Generic[I, O]):
 
     endpoint: str
     service: str
-    operation: Union[nexusrpc.Operation[I, O], str]
+    operation: Union[
+        nexusrpc.interface.Operation[I, O],
+        Callable[[Any], nexusrpc.handler.Operation[I, O]],
+        str,
+    ]
     input: I
     schedule_to_close_timeout: Optional[timedelta]
     headers: Optional[Mapping[str, str]]
 
+    # Cached properties, initialized in __post_init__
+    _operation_name: str = field(init=False, repr=False)
+    _input_type: Optional[Type[I]] = field(init=False, repr=False)
+    _output_type: Optional[Type[O]] = field(init=False, repr=False)
+
+    def __post_init__(self) -> None:
+        if isinstance(self.operation, str):
+            self._operation_name = self.operation
+            self._input_type = None
+            self._output_type = None
+        elif isinstance(self.operation, nexusrpc.interface.Operation):
+            self._operation_name = self.operation.name
+            self._input_type = self.operation.input_type
+            self._output_type = self.operation.output_type
+        elif isinstance(self.operation, Callable):
+            defn = getattr(self.operation, "__nexus_operation__", None)
+            if isinstance(defn, nexusrpc.handler.NexusOperationDefinition):
+                self._operation_name = defn.name
+                self._input_type = defn.input_type
+                self._output_type = defn.output_type
+            else:
+                raise ValueError(
+                    f"Operation callable is not a Nexus operation: {self.operation}"
+                )
+        else:
+            raise ValueError(f"Operation is not a Nexus operation: {self.operation}")
+
     @property
     def operation_name(self) -> str:
-        return (
-            self.operation.name
-            if isinstance(self.operation, nexusrpc.Operation)
-            else self.operation
-        )
+        return self._operation_name
 
     @property
     def input_type(self) -> Optional[Type[I]]:
-        return (
-            self.operation.input_type
-            if isinstance(self.operation, nexusrpc.Operation)
-            else None
-        )
+        return self._input_type
 
     @property
     def output_type(self) -> Optional[Type[O]]:
-        return (
-            self.operation.output_type
-            if isinstance(self.operation, nexusrpc.Operation)
-            else None
-        )
+        return self._output_type
 
 
 @dataclass

temporalio/worker/_workflow_instance.py

@@ -1559,7 +1559,11 @@ async def workflow_start_nexus_operation(
         self,
         endpoint: str,
         service: str,
-        operation: Union[nexusrpc.Operation[I, O], str],
+        operation: Union[
+            nexusrpc.interface.Operation[I, O],
+            Callable[[Any], nexusrpc.handler.Operation[I, O]],
+            str,
+        ],
         input: Any,
         schedule_to_close_timeout: Optional[timedelta] = None,
         headers: Optional[Mapping[str, str]] = None,

temporalio/workflow.py

@@ -840,7 +840,11 @@ async def workflow_start_nexus_operation(
         self,
         endpoint: str,
         service: str,
-        operation: Union[nexusrpc.Operation[I, O], str],
+        operation: Union[
+            nexusrpc.interface.Operation[I, O],
+            Callable[[Any], nexusrpc.handler.Operation[I, O]],
+            str,
+        ],
         input: Any,
         schedule_to_close_timeout: Optional[timedelta] = None,
         headers: Optional[Mapping[str, str]] = None,
@@ -4385,7 +4389,11 @@ def operation_token(self) -> Optional[str]:
 async def start_nexus_operation(
     endpoint: str,
     service: str,
-    operation: Union[nexusrpc.Operation[I, O], str],
+    operation: Union[
+        nexusrpc.interface.Operation[I, O],
+        Callable[[Any], nexusrpc.handler.Operation[I, O]],
+        str,
+    ],
     input: Any,
     *,
     schedule_to_close_timeout: Optional[timedelta] = None,
@@ -5158,7 +5166,11 @@ def __init__(
     # TODO(dan): overloads: no-input, operation name, ret type
     async def start_operation(
         self,
-        operation: Union[nexusrpc.interface.Operation[I, O], str],
+        operation: Union[
+            nexusrpc.interface.Operation[I, O],
+            Callable[[Any], nexusrpc.handler.Operation[I, O]],
+            str,
+        ],
         input: I,
         *,
         schedule_to_close_timeout: Optional[timedelta] = None,
@@ -5180,7 +5192,7 @@ async def execute_operation(
         self,
         operation: Union[
             nexusrpc.interface.Operation[I, O],
-            Callable[..., nexusrpc.handler.Operation[I, O]],
+            Callable[[Any], nexusrpc.handler.Operation[I, O]],
             str,
         ],
         input: I,