RTU: Copy operation factory getter/setter from nexusrpc

Author: dandavison

temporalio/nexus/_util.py

@@ -125,6 +125,40 @@ def get_callable_name(fn: Callable[..., Any]) -> str:
     return method_name
 
 
+# TODO(nexus-preview) Copied from nexusrpc
+def get_operation_factory(
+    obj: Any,
+) -> tuple[
+    Optional[Callable[[Any], OperationHandler[InputT, OutputT]]],
+    Optional[nexusrpc.Operation[InputT, OutputT]],
+]:
+    """Return the :py:class:`Operation` for the object along with the factory function.
+
+    ``obj`` should be a decorated operation start method.
+    """
+    op_defn = get_operation_definition(obj)
+    if op_defn:
+        factory = obj
+    else:
+        if factory := getattr(obj, "__nexus_operation_factory__", None):
+            op_defn = get_operation_definition(factory)
+    if not isinstance(op_defn, nexusrpc.Operation):
+        return None, None
+    return factory, op_defn
+
+
+# TODO(nexus-preview) Copied from nexusrpc
+def set_operation_factory(
+    obj: Any,
+    operation_factory: Callable[[Any], OperationHandler[InputT, OutputT]],
+) -> None:
+    """Set the :py:class:`OperationHandler` factory for this object.
+
+    ``obj`` should be an operation start method.
+    """
+    setattr(obj, "__nexus_operation_factory__", operation_factory)
+
+
 # Copied from https://github.com/modelcontextprotocol/python-sdk
 #
 # Copyright (c) 2024 Anthropic, PBC.

temporalio/worker/_interceptor.py

@@ -25,6 +25,7 @@
 import temporalio.api.common.v1
 import temporalio.common
 import temporalio.nexus
+import temporalio.nexus._util
 import temporalio.workflow
 from temporalio.workflow import VersioningIntent
 
@@ -313,7 +314,7 @@ def __post_init__(self) -> None:
             self._operation_name = self.operation
             self._input_type = None
         elif isinstance(self.operation, Callable):
-            _, op = nexusrpc.get_operation_factory(self.operation)
+            _, op = temporalio.nexus._util.get_operation_factory(self.operation)
             if isinstance(op, nexusrpc.Operation):
                 self._operation_name = op.name
                 self._input_type = op.input_type

tests/nexus/test_dynamic_creation_of_user_handler_classes.py

@@ -3,10 +3,10 @@
 import httpx
 import nexusrpc.handler
 import pytest
-from nexusrpc import get_operation_factory
 from nexusrpc.handler import sync_operation
 
 from temporalio.client import Client
+from temporalio.nexus._util import get_operation_factory
 from temporalio.worker import Worker
 from tests.helpers.nexus import create_nexus_endpoint
 

tests/nexus/test_handler_operation_definitions.py

@@ -11,6 +11,7 @@
 
 from temporalio import nexus
 from temporalio.nexus import WorkflowRunOperationContext, workflow_run_operation
+from temporalio.nexus._util import get_operation_factory
 
 
 @dataclass
@@ -92,9 +93,7 @@ async def test_collected_operation_names(
     assert isinstance(service_defn, nexusrpc.ServiceDefinition)
     assert service_defn.name == "Service"
     for method_name, expected_op in test_case.expected_operations.items():
-        _, actual_op = nexusrpc.get_operation_factory(
-            getattr(test_case.Service, method_name)
-        )
+        _, actual_op = get_operation_factory(getattr(test_case.Service, method_name))
         assert isinstance(actual_op, nexusrpc.Operation)
         assert actual_op.name == expected_op.name
         assert actual_op.input_type == expected_op.input_type