Clean up existing code

Author: dandavison

temporalio/worker/_interceptor.py

@@ -388,23 +388,23 @@ async def signal_external_workflow(
 
     def start_activity(
         self, input: StartActivityInput
-    ) -> temporalio.workflow.ActivityHandle:
+    ) -> temporalio.workflow.ActivityHandle[Any]:
         """Called for every :py:func:`temporalio.workflow.start_activity` and
         :py:func:`temporalio.workflow.execute_activity` call.
         """
         return self.next.start_activity(input)
 
     async def start_child_workflow(
         self, input: StartChildWorkflowInput
-    ) -> temporalio.workflow.ChildWorkflowHandle:
+    ) -> temporalio.workflow.ChildWorkflowHandle[Any, Any]:
         """Called for every :py:func:`temporalio.workflow.start_child_workflow`
         and :py:func:`temporalio.workflow.execute_child_workflow` call.
         """
         return await self.next.start_child_workflow(input)
 
     def start_local_activity(
         self, input: StartLocalActivityInput
-    ) -> temporalio.workflow.ActivityHandle:
+    ) -> temporalio.workflow.ActivityHandle[Any]:
         """Called for every :py:func:`temporalio.workflow.start_local_activity`
         and :py:func:`temporalio.workflow.execute_local_activity` call.
         """

temporalio/worker/_workflow_instance.py

@@ -1537,25 +1537,26 @@ def _outbound_schedule_activity(
         self,
         input: Union[StartActivityInput, StartLocalActivityInput],
     ) -> _ActivityHandle:
+        # A ScheduleActivityTask command always results in an ActivityTaskScheduled event,
+        # so this function returns the handle immediately. This is similar to nexus
+        # operation but differs from child workflow.
+
         # Validate
         if not input.start_to_close_timeout and not input.schedule_to_close_timeout:
             raise ValueError(
                 "Activity must have start_to_close_timeout or schedule_to_close_timeout"
             )
 
-        handle: Optional[_ActivityHandle] = None
+        handle: _ActivityHandle
 
         # Function that runs in the handle
         async def run_activity() -> Any:
-            nonlocal handle
-            assert handle
             while True:
                 # Mark it as started each loop because backoff could cause it to
                 # be marked as unstarted
                 handle._started = True
                 try:
-                    # We have to shield because we don't want the underlying
-                    # result future to be cancelled
+                    # Shield so that future itself is not cancelled
                     return await asyncio.shield(handle._result_fut)
                 except _ActivityDoBackoffError as err:
                     # We have to sleep then reschedule. Note this sleep can be
@@ -1615,12 +1616,16 @@ async def _outbound_signal_external_workflow(
     async def _outbound_start_child_workflow(
         self, input: StartChildWorkflowInput
     ) -> _ChildWorkflowHandle:
-        handle: Optional[_ChildWorkflowHandle] = None
+        # A StartChildWorkflowExecution command results in a
+        # StartChildWorkflowExecutionInitiated event, but the start may fail (e.g. due to
+        # workflow ID collision). Therefore this function does not return the handle until
+        # a future activation contains an event indicating start success / failure. This
+        # differs from activity and nexus operation.
+
+        handle: _ChildWorkflowHandle
 
         # Common code for handling cancel for start and run
         def apply_child_cancel_error() -> None:
-            nonlocal handle
-            assert handle
             # Send a cancel request to the child
             cancel_command = self._add_command()
             handle._apply_cancel_command(cancel_command)
@@ -1638,12 +1643,9 @@ def apply_child_cancel_error() -> None:
 
         # Function that runs in the handle
         async def run_child() -> Any:
-            nonlocal handle
             while True:
-                assert handle
                 try:
-                    # We have to shield because we don't want the future itself
-                    # to be cancelled
+                    # Shield so that future itself is not cancelled
                     return await asyncio.shield(handle._result_fut)
                 except asyncio.CancelledError:
                     apply_child_cancel_error()
@@ -1658,8 +1660,7 @@ async def run_child() -> Any:
         # Wait on start before returning
         while True:
             try:
-                # We have to shield because we don't want the future itself
-                # to be cancelled
+                # Shield so that future itself is not cancelled
                 await asyncio.shield(handle._start_fut)
                 return handle
             except asyncio.CancelledError:
@@ -2391,17 +2392,17 @@ async def signal_external_workflow(
 
     def start_activity(
         self, input: StartActivityInput
-    ) -> temporalio.workflow.ActivityHandle:
+    ) -> temporalio.workflow.ActivityHandle[Any]:
         return self._instance._outbound_schedule_activity(input)
 
     async def start_child_workflow(
         self, input: StartChildWorkflowInput
-    ) -> temporalio.workflow.ChildWorkflowHandle:
+    ) -> temporalio.workflow.ChildWorkflowHandle[Any, Any]:
         return await self._instance._outbound_start_child_workflow(input)
 
     def start_local_activity(
         self, input: StartLocalActivityInput
-    ) -> temporalio.workflow.ActivityHandle:
+    ) -> temporalio.workflow.ActivityHandle[Any]:
         return self._instance._outbound_schedule_activity(input)
 
 

tests/helpers/__init__.py

@@ -4,7 +4,7 @@
 import uuid
 from contextlib import closing
 from datetime import timedelta
-from typing import Any, Awaitable, Callable, Optional, Sequence, Type, TypeVar, Union
+from typing import Any, Awaitable, Callable, Optional, Sequence, Type, TypeVar
 
 from temporalio.api.common.v1 import WorkflowExecution
 from temporalio.api.enums.v1 import IndexedValueType