Rearrange activity handle classes

Author: dandavison

temporalio/client.py

@@ -2849,134 +2849,6 @@ async def workflow_handle(self) -> WorkflowHandle[SelfType, ReturnType]:
         return await self._workflow_handle
 
 
-@dataclass(frozen=True)
-# TODO: in the future when messages can be sent to activities, we will want the activity handle to
-# be generic in the activity type in addition to the return type (as WorkflowHandle), to support
-# static type inference for signal/query/update.
-class ActivityHandle(Generic[ReturnType]):
-    """Handle representing a activity execution."""
-
-    client: Client
-    """The client that created the activity handle."""
-
-    id: str
-    """The activity ID."""
-
-    run_id: Optional[str] = None
-    """The run ID of the activity."""
-
-    # TODO: do we support something like `follow_runs: bool`?
-    async def result(
-        self,
-        *,
-        rpc_metadata: Mapping[str, Union[str, bytes]] = {},
-        rpc_timeout: Optional[timedelta] = None,
-    ) -> ReturnType:
-        """Wait for result of the activity.
-
-        Args:
-            rpc_metadata: Headers used on the RPC call. Keys here override
-                client-level RPC metadata keys.
-            rpc_timeout: Optional RPC deadline to set for each RPC call. Note,
-                this is the timeout for each history RPC call not this overall
-                function.
-
-        Returns:
-            The result of the activity.
-
-        Raises:
-            :py:class:`ActivityFailureError`: If the activity completed with a failure.
-        """
-        # Repeatedly issues workflowservice GetActivityResult long-polls.
-        raise NotImplementedError
-
-    async def cancel(
-        self,
-        *,
-        reason: Optional[str] = None,
-        rpc_metadata: Mapping[str, Union[str, bytes]] = {},
-        rpc_timeout: Optional[timedelta] = None,
-    ) -> None:
-        """Request cancellation of the activity.
-
-        Requesting cancellation of an activity does not automatically transition the activity to
-        canceled status. If the activity is heartbeating, a :py:class:`exceptions.CancelledError`
-        exception will be raised when receiving the heartbeat response; if the activity allows this
-        exception to bubble out, the activity will transition to canceled status. If the activity it
-        is not heartbeating, this method will have no effect on activity status.
-
-        Args:
-            reason: Reason for the cancellation. Recorded and available via describe.
-            rpc_metadata: Headers used on the RPC call.
-            rpc_timeout: Optional RPC deadline to set for the RPC call.
-        """
-        raise NotImplementedError
-
-    async def complete(
-        self,
-        result: Optional[Any] = temporalio.common._arg_unset,
-        *,
-        rpc_metadata: Mapping[str, Union[str, bytes]] = {},
-        rpc_timeout: Optional[timedelta] = None,
-    ) -> None:
-        """Complete the activity.
-
-        Args:
-            result: Result of the activity if any.
-            rpc_metadata: Headers used on the RPC call. Keys here override
-                client-level RPC metadata keys.
-            rpc_timeout: Optional RPC deadline to set for the RPC call.
-        """
-        await self.client._impl.complete_async_activity(
-            CompleteAsyncActivityInput(
-                id_or_token=ActivityIDReference(
-                    activity_id=self.id,
-                    run_id=self.run_id,
-                ),
-                result=result,
-                rpc_metadata=rpc_metadata,
-                rpc_timeout=rpc_timeout,
-            ),
-        )
-
-    async def terminate(
-        self,
-        *,
-        reason: Optional[str] = None,
-        rpc_metadata: Mapping[str, Union[str, bytes]] = {},
-        rpc_timeout: Optional[timedelta] = None,
-    ) -> None:
-        """Terminate the activity execution immediately.
-
-        Termination does not reach the worker and the activity code cannot react to it.
-        A terminated activity may have a running attempt and will be requested to be
-        canceled by the server when it heartbeats.
-
-        Args:
-            reason: Reason for the termination.
-            rpc_metadata: Headers used on the RPC call.
-            rpc_timeout: Optional RPC deadline to set for the RPC call.
-        """
-        raise NotImplementedError
-
-    async def describe(
-        self,
-        *,
-        rpc_metadata: Mapping[str, Union[str, bytes]] = {},
-        rpc_timeout: Optional[timedelta] = None,
-    ) -> ActivityExecutionDescription:
-        """Describe the activity execution.
-
-        Args:
-            rpc_metadata: Headers used on the RPC call.
-            rpc_timeout: Optional RPC deadline to set for the RPC call.
-
-        Returns:
-            Activity execution description.
-        """
-        raise NotImplementedError
-
-
 class ActivityExecutionAsyncIterator:
     """Asynchronous iterator for :py:class:`ActivityExecution` values."""
 
@@ -3124,15 +2996,11 @@ class AsyncActivityIDReference:
     activity_id: str
 
 
-class AsyncActivityHandle:
-    """Handle representing an external activity for completion and heartbeat."""
+class _BaseActivityHandle:
+    """Handle representing an activity."""
 
-    def __init__(
-        self, client: Client, id_or_token: Union[AsyncActivityIDReference, bytes]
-    ) -> None:
-        """Create an async activity handle."""
-        self._client = client
-        self._id_or_token = id_or_token
+    _client: Client
+    _id_or_token: Union[ActivityIDReference, AsyncActivityIDReference, bytes]
 
     async def heartbeat(
         self,
@@ -3232,6 +3100,127 @@ async def report_cancellation(
         )
 
 
+# TODO: This name is suboptimal now. We could deprecate it and introduce WorkflowActivityHandle as a
+# preferred alias.
+class AsyncActivityHandle(_BaseActivityHandle):
+    """Handle representing an activity started by a workflow."""
+
+    def __init__(
+        self, client: Client, id_or_token: Union[AsyncActivityIDReference, bytes]
+    ) -> None:
+        """Create an handle to an activity started by a workflow."""
+        self._client = client
+        self._id_or_token = id_or_token
+
+
+# TODO: in the future when messages can be sent to activities, we will want the activity handle to
+# be generic in the activity type in addition to the return type (as WorkflowHandle), to support
+# static type inference for signal/query/update.
+class ActivityHandle(Generic[ReturnType], _BaseActivityHandle):
+    """Handle representing a standalone activity execution."""
+
+    def __init__(self, client: Client, id: str, run_id: Optional[str] = None) -> None:
+        """Create an activity handle.
+
+        Args:
+            client: The client that created the activity handle.
+            id: The activity ID.
+            run_id: The run ID of the activity.
+        """
+        self._client = client
+        self._id_or_token = ActivityIDReference(activity_id=id, run_id=run_id)
+        self.run_id = run_id
+
+    # TODO: do we support something like `follow_runs: bool`?
+    async def result(
+        self,
+        *,
+        rpc_metadata: Mapping[str, Union[str, bytes]] = {},
+        rpc_timeout: Optional[timedelta] = None,
+    ) -> ReturnType:
+        """Wait for result of the activity.
+
+        Args:
+            rpc_metadata: Headers used on the RPC call. Keys here override
+                client-level RPC metadata keys.
+            rpc_timeout: Optional RPC deadline to set for each RPC call. Note,
+                this is the timeout for each history RPC call not this overall
+                function.
+
+        Returns:
+            The result of the activity.
+
+        Raises:
+            :py:class:`ActivityFailureError`: If the activity completed with a failure.
+        """
+        # Repeatedly issues workflowservice GetActivityResult long-polls.
+        raise NotImplementedError
+
+    async def cancel(
+        self,
+        *,
+        reason: Optional[str] = None,
+        rpc_metadata: Mapping[str, Union[str, bytes]] = {},
+        rpc_timeout: Optional[timedelta] = None,
+    ) -> None:
+        """Request cancellation of the activity.
+
+        Requesting cancellation of an activity does not automatically transition the activity to
+        canceled status. If the activity is heartbeating, a :py:class:`exceptions.CancelledError`
+        exception will be raised when receiving the heartbeat response; if the activity allows this
+        exception to bubble out, the activity will transition to canceled status. If the activity it
+        is not heartbeating, this method will have no effect on activity status.
+
+        Args:
+            reason: Reason for the cancellation. Recorded and available via describe.
+            rpc_metadata: Headers used on the RPC call.
+            rpc_timeout: Optional RPC deadline to set for the RPC call.
+        """
+        raise NotImplementedError
+
+    async def terminate(
+        self,
+        *,
+        reason: Optional[str] = None,
+        rpc_metadata: Mapping[str, Union[str, bytes]] = {},
+        rpc_timeout: Optional[timedelta] = None,
+    ) -> None:
+        """Terminate the activity execution immediately.
+
+        Termination does not reach the worker and the activity code cannot react to it.
+        A terminated activity may have a running attempt and will be requested to be
+        canceled by the server when it heartbeats.
+
+        Args:
+            reason: Reason for the termination.
+            rpc_metadata: Headers used on the RPC call.
+            rpc_timeout: Optional RPC deadline to set for the RPC call.
+        """
+        raise NotImplementedError
+
+    async def describe(
+        self,
+        *,
+        rpc_metadata: Mapping[str, Union[str, bytes]] = {},
+        rpc_timeout: Optional[timedelta] = None,
+    ) -> ActivityExecutionDescription:
+        """Describe the activity execution.
+
+        Args:
+            rpc_metadata: Headers used on the RPC call.
+            rpc_timeout: Optional RPC deadline to set for the RPC call.
+
+        Returns:
+            Activity execution description.
+        """
+        raise NotImplementedError
+
+    # TODO:
+    # pause
+    # reset
+    # update_options
+
+
 @dataclass
 class WorkflowExecution:
     """Info for a single workflow execution run."""
@@ -5867,7 +5856,7 @@ class StartWorkflowUpdateWithStartInput:
 class HeartbeatAsyncActivityInput:
     """Input for :py:meth:`OutboundInterceptor.heartbeat_async_activity`."""
 
-    id_or_token: Union[AsyncActivityIDReference, bytes]
+    id_or_token: Union[AsyncActivityIDReference, ActivityIDReference, bytes]
     details: Sequence[Any]
     rpc_metadata: Mapping[str, Union[str, bytes]]
     rpc_timeout: Optional[timedelta]