Swap type aliases back to classes to avoid printing or runtime type checking concerns

Author: VegetarianOrc

temporalio/nexus/_operation_context.py

@@ -19,7 +19,6 @@
     Any,
     Concatenate,
     Generic,
-    TypeAlias,
     TypeVar,
     overload,
 )
@@ -29,6 +28,7 @@
     OperationContext,
     StartOperationContext,
 )
+from typing_extensions import Self
 
 import temporalio.api.common.v1
 import temporalio.api.workflowservice.v1
@@ -549,8 +549,32 @@ def set(self) -> None:
         _temporal_cancel_operation_context.set(self)
 
 
-TemporalNexusStartOperationContext: TypeAlias = StartOperationContext
-TemporalNexusCancelOperationContext: TypeAlias = CancelOperationContext
+class TemporalNexusStartOperationContext(StartOperationContext):
+    """Context received by a Temporal Nexus operation when it is started.
+
+    .. warning::
+       This API is experimental and unstable.
+    """
+
+    @classmethod
+    def _from_start_operation_context(cls, ctx: StartOperationContext) -> Self:
+        return cls(
+            **{f.name: getattr(ctx, f.name) for f in dataclasses.fields(ctx)},
+        )
+
+
+class TemporalNexusCancelOperationContext(CancelOperationContext):
+    """Context received by a Temporal Nexus operation when it is canceled.
+
+    .. warning::
+       This API is experimental and unstable.
+    """
+
+    @classmethod
+    def _from_cancel_operation_context(cls, ctx: CancelOperationContext) -> Self:
+        return cls(
+            **{f.name: getattr(ctx, f.name) for f in dataclasses.fields(ctx)},
+        )
 
 
 class LoggerAdapter(logging.LoggerAdapter):

temporalio/nexus/_operation_handlers.py

@@ -150,7 +150,10 @@ async def start(
            This API is experimental and unstable.
         """
         nexus_client = _TemporalNexusClient()
-        result = await self.start_operation(ctx, nexus_client, input)
+        start_ctx = TemporalNexusStartOperationContext._from_start_operation_context(
+            ctx
+        )
+        result = await self.start_operation(start_ctx, nexus_client, input)
         return result._to_nexus_result()
 
     async def cancel(self, ctx: CancelOperationContext, token: str) -> None:
@@ -167,9 +170,12 @@ async def cancel(self, ctx: CancelOperationContext, token: str) -> None:
                 type=HandlerErrorType.INTERNAL,
             ) from err
 
+        cancel_ctx = TemporalNexusCancelOperationContext._from_cancel_operation_context(
+            ctx
+        )
         match operation_token.type:
             case OperationTokenType.WORKFLOW:
-                await self.cancel_workflow_run(ctx, operation_token.workflow_id)
+                await self.cancel_workflow_run(cancel_ctx, operation_token.workflow_id)
 
     async def cancel_workflow_run(
         self,

temporalio/nexus/_util.py

@@ -5,7 +5,6 @@
 import typing
 import warnings
 from collections.abc import Awaitable, Callable
-from dataclasses import dataclass
 from typing import (
     Any,
 )
@@ -31,16 +30,6 @@
 )
 
 
-@dataclass(frozen=True)
-class _ExpectedParamType:
-    runtime_type: type[Any]
-    display_name: str | None = None
-
-    @property
-    def name(self) -> str:
-        return self.display_name or self.runtime_type.__name__
-
-
 def get_workflow_run_start_method_input_and_output_type_annotations(
     start: Callable[
         [NexusServiceType, WorkflowRunOperationContext, InputT],
@@ -57,7 +46,7 @@ def get_workflow_run_start_method_input_and_output_type_annotations(
     """
     return _get_wrapped_start_method_input_and_output_type_annotations(
         start,
-        expected_param_types=(_ExpectedParamType(WorkflowRunOperationContext),),
+        expected_param_types=(WorkflowRunOperationContext,),
         expected_return_origin=WorkflowHandle,
     )
 
@@ -84,11 +73,8 @@ def get_temporal_operation_start_method_input_and_output_type_annotations(
     return _get_wrapped_start_method_input_and_output_type_annotations(
         start,
         expected_param_types=(
-            _ExpectedParamType(
-                TemporalNexusStartOperationContext,
-                "TemporalNexusStartOperationContext",
-            ),
-            _ExpectedParamType(TemporalNexusClient),
+            TemporalNexusStartOperationContext,
+            TemporalNexusClient,
         ),
         expected_return_origin=TemporalOperationResult,
     )
@@ -97,7 +83,7 @@ def get_temporal_operation_start_method_input_and_output_type_annotations(
 def _get_wrapped_start_method_input_and_output_type_annotations(
     start: Callable[..., Any],
     *,
-    expected_param_types: tuple[_ExpectedParamType, ...],
+    expected_param_types: tuple[type[Any], ...],
     expected_return_origin: type[Any],
 ) -> tuple[
     type[Any] | None,
@@ -135,7 +121,7 @@ def _get_wrapped_start_method_input_and_output_type_annotations(
 def _get_start_method_input_and_output_type_annotations(
     start: Callable[..., Any],
     *,
-    expected_param_types: tuple[_ExpectedParamType, ...],
+    expected_param_types: tuple[type[Any], ...],
 ) -> tuple[
     type[Any] | None,
     type[Any] | None,
@@ -164,10 +150,10 @@ def _get_start_method_input_and_output_type_annotations(
         for index, (param_type, expected_param_type) in enumerate(
             zip(param_types, expected_param_types), start=1
         ):
-            if not issubclass(expected_param_type.runtime_type, param_type):
+            if not issubclass(expected_param_type, param_type):
                 warnings.warn(
                     f"Expected parameter {index} of {start} to be an instance of "
-                    f"{expected_param_type.name}, but is {param_type}."
+                    f"{expected_param_type.__name__}, but is {param_type}."
                 )
                 input_type = None