Return self in AsyncActivityHandle.with_handle

Author: dandavison

temporalio/client.py

@@ -41,7 +41,7 @@
 import google.protobuf.json_format
 import google.protobuf.timestamp_pb2
 from google.protobuf.internal.containers import MessageMap
-from typing_extensions import Concatenate, Required, TypedDict
+from typing_extensions import Concatenate, Required, Self, TypedDict
 
 import temporalio.api.common.v1
 import temporalio.api.enums.v1
@@ -2852,17 +2852,24 @@ async def report_cancellation(
             ),
         )
 
-    # TODO(dan): should this return Self (requiring that the user's subclass has the same
-    # constructor signature)? CompositePayloadConverter.with_context does.
-    def with_context(self, context: SerializationContext) -> AsyncActivityHandle:
+    def with_context(self, context: SerializationContext) -> Self:
         """Create a new AsyncActivityHandle with a different serialization context.
 
         Payloads received by the activity will be decoded and deserialized using a data converter
         with :py:class:`ActivitySerializationContext` set as context. If you are using a custom data
         converter that makes use of this context then you can use this method to supply matching
         context data to the data converter used to serialize and encode the outbound payloads.
         """
-        return AsyncActivityHandle(
+        data_converter = self._client.data_converter.with_context(context)
+        if data_converter == self._client.data_converter:
+            return self
+        cls = type(self)
+        if cls.__init__ is not AsyncActivityHandle.__init__:
+            raise TypeError(
+                "If you have subclassed AsyncActivityHandle and overridden the __init__ method "
+                "then you must override with_context to return an instance of your class."
+            )
+        return cls(
             self._client,
             self._id_or_token,
             self._client.data_converter.with_context(context),

tests/test_serialization_context.py

@@ -24,7 +24,12 @@
 import temporalio.api.common.v1
 import temporalio.api.failure.v1
 from temporalio import activity, workflow
-from temporalio.client import Client, WorkflowFailureError, WorkflowUpdateFailedError
+from temporalio.client import (
+    AsyncActivityHandle,
+    Client,
+    WorkflowFailureError,
+    WorkflowUpdateFailedError,
+)
 from temporalio.common import RetryPolicy
 from temporalio.contrib.pydantic import PydanticJSONPlainPayloadConverter
 from temporalio.converter import (
@@ -610,6 +615,60 @@ async def test_async_activity_completion_payload_conversion(
         ]
 
 
+class MyAsyncActivityHandle(AsyncActivityHandle):
+    def my_method(self) -> None:
+        pass
+
+
+class MyAsyncActivityHandleWithOverriddenConstructor(AsyncActivityHandle):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+
+    def my_method(self) -> None:
+        pass
+
+
+def test_subclassed_async_activity_handle(client: Client):
+    activity_context = ActivitySerializationContext(
+        namespace="default",
+        workflow_id="workflow-id",
+        workflow_type="workflow-type",
+        activity_type="activity-type",
+        activity_task_queue="activity-task-queue",
+        is_local=False,
+    )
+    handle = MyAsyncActivityHandle(client=client, id_or_token=b"task-token")
+    # This works because the data converter does not use context so AsyncActivityHandle.with_context
+    # returns self
+    assert isinstance(handle.with_context(activity_context), MyAsyncActivityHandle)
+
+    # This time the data converter uses context so AsyncActivityHandle.with_context attempts to
+    # return a new instance of the user's subclass. It works, because they have not overridden the
+    # constructor.
+    client_config = client.config()
+    client_config["data_converter"] = dataclasses.replace(
+        DataConverter.default,
+        payload_converter_class=SerializationContextCompositePayloadConverter,
+    )
+    client = Client(**client_config)
+    handle = MyAsyncActivityHandle(client=client, id_or_token=b"task-token")
+    assert isinstance(handle.with_context(activity_context), MyAsyncActivityHandle)
+
+    # Finally, a user attempts the same but having overridden the constructor. This fails:
+    # AsyncActivityHandle.with_context refuses to attempt to create an instance of their subclass.
+    handle2 = MyAsyncActivityHandleWithOverriddenConstructor(
+        client=client, id_or_token=b"task-token"
+    )
+    with pytest.raises(
+        TypeError,
+        match="you must override with_context to return an instance of your class",
+    ):
+        assert isinstance(
+            handle2.with_context(activity_context),
+            MyAsyncActivityHandleWithOverriddenConstructor,
+        )
+
+
 # Signal test