use object reference instead of function, picklable

THardy98 · THardy98 · commit 4451e77bffee · 2025-05-03T14:07:30.000-04:00
diff --git a/temporalio/activity.py b/temporalio/activity.py
@@ -174,7 +174,7 @@ class _Context:
         temporalio.converter.PayloadConverter,
     ]
     runtime_metric_meter: Optional[temporalio.common.MetricMeter]
-    cancellation_details: Callable[[], Optional[ActivityCancellationDetails]]
+    cancellation_details: Optional[ActivityCancellationDetails] = None
     _logger_details: Optional[Mapping[str, Any]] = None
     _payload_converter: Optional[temporalio.converter.PayloadConverter] = None
     _metric_meter: Optional[temporalio.common.MetricMeter] = None
@@ -289,7 +289,7 @@ def info() -> Info:
 
 def cancellation_details() -> Optional[ActivityCancellationDetails]:
     """Cancellation details of the currenct activity, if any"""
-    return _Context.current().cancellation_details()
+    return _Context.current().cancellation_details
 
 
 def heartbeat(*details: Any) -> None:
diff --git a/temporalio/testing/_activity.py b/temporalio/testing/_activity.py
@@ -168,7 +168,7 @@ def __init__(
             else self.cancel_thread_raiser.shielded,
             payload_converter_class_or_instance=env.payload_converter,
             runtime_metric_meter=env.metric_meter,
-            cancellation_details=lambda: env.cancellation_details,
+            cancellation_details=env.cancellation_details,
         )
         self.task: Optional[asyncio.Task] = None
 
diff --git a/temporalio/worker/_activity.py b/temporalio/worker/_activity.py
@@ -15,7 +15,7 @@
 import warnings
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from datetime import datetime, timedelta, timezone
 from typing import (
     Any,
@@ -214,9 +214,8 @@ def _cancel(
             warnings.warn(f"Cannot find activity to cancel for token {task_token!r}")
             return
         logger.debug("Cancelling activity %s, reason: %s", task_token, cancel.reason)
-        activity.cancellation_details = (
-            temporalio.activity.ActivityCancellationDetails._fromProto(cancel.details)
-        )
+        activity.cancellation_details.cancelled = cancel.details.is_cancelled
+        activity.cancellation_details.paused = cancel.details.is_paused
         activity.cancel(cancelled_by_request=True)
 
     def _heartbeat(self, task_token: bytes, *details: Any) -> None:
@@ -424,8 +423,6 @@ async def _run_activity(
                 headers=start.header_fields,
             )
 
-            running_activity.cancellation_details = None
-
             # Set the context early so the logging adapter works and
             # interceptors have it
             temporalio.activity._Context.set(
@@ -441,8 +438,7 @@ async def _run_activity(
                     runtime_metric_meter=None
                     if sync_non_threaded
                     else self._metric_meter,
-                    # Function reference to the running activity's cancellation details
-                    cancellation_details=lambda: running_activity.cancellation_details,
+                    cancellation_details=running_activity.cancellation_details,
                 )
             )
             temporalio.activity.logger.debug("Starting activity")
@@ -553,8 +549,8 @@ class _RunningActivity:
     done: bool = False
     cancelled_by_request: bool = False
     cancelled_due_to_heartbeat_error: Optional[Exception] = None
-    cancellation_details: Optional[temporalio.activity.ActivityCancellationDetails] = (
-        None
+    cancellation_details: temporalio.activity.ActivityCancellationDetails = field(
+        default_factory=temporalio.activity.ActivityCancellationDetails
     )
 
     def cancel(
@@ -751,9 +747,7 @@ def _execute_sync_activity(
         temporalio.converter.PayloadConverter,
     ],
     runtime_metric_meter: Optional[temporalio.common.MetricMeter],
-    cancellation_details: Callable[
-        [], Optional[temporalio.activity.ActivityCancellationDetails]
-    ],
+    cancellation_details: Optional[temporalio.activity.ActivityCancellationDetails],
     fn: Callable[..., Any],
     *args: Any,
 ) -> Any:

Original file line number	Diff line number	Diff line change
`@@ -168,7 +168,7 @@ def __init__(`
`168`	`168`	`else self.cancel_thread_raiser.shielded,`
`169`	`169`	`payload_converter_class_or_instance=env.payload_converter,`
`170`	`170`	`runtime_metric_meter=env.metric_meter,`
`171`		`- cancellation_details=lambda: env.cancellation_details,`
	`171`	`+ cancellation_details=env.cancellation_details,`
`172`	`172`	`)`
`173`	`173`	`self.task: Optional[asyncio.Task] = None`
`174`	`174`