Skip to content

Commit 53e0fa8

Browse files
brianstrauchbrianstrauch
authored
Merge branch 'main' into dependabot/uv/langsmith-0.8.0
2 parents 6c51ddf + 43a3c16 commit 53e0fa8

19 files changed

Lines changed: 2358 additions & 191 deletions

‎temporalio/client/_nexus.py‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -611,6 +611,35 @@ async def start_operation(
611611
rpc_timeout: timedelta | None = None,
612612
) -> NexusOperationHandle[OutputT]: ...
613613

614+
# Overload for temporal_operation methods
615+
@overload
616+
@abstractmethod
617+
async def start_operation(
618+
self,
619+
operation: Callable[
620+
[
621+
NexusServiceType,
622+
temporalio.nexus.TemporalNexusStartOperationContext,
623+
temporalio.nexus.TemporalNexusClient,
624+
InputT,
625+
],
626+
Awaitable[temporalio.nexus.TemporalOperationResult[OutputT]],
627+
],
628+
arg: InputT,
629+
*,
630+
id: str,
631+
id_reuse_policy: temporalio.common.NexusOperationIDReusePolicy = temporalio.common.NexusOperationIDReusePolicy.ALLOW_DUPLICATE,
632+
id_conflict_policy: temporalio.common.NexusOperationIDConflictPolicy = temporalio.common.NexusOperationIDConflictPolicy.FAIL,
633+
schedule_to_close_timeout: timedelta | None = None,
634+
schedule_to_start_timeout: timedelta | None = None,
635+
start_to_close_timeout: timedelta | None = None,
636+
search_attributes: temporalio.common.TypedSearchAttributes | None = None,
637+
summary: str | None = None,
638+
headers: Mapping[str, str] | None = None,
639+
rpc_metadata: Mapping[str, str | bytes] = {},
640+
rpc_timeout: timedelta | None = None,
641+
) -> NexusOperationHandle[OutputT]: ...
642+
614643
@abstractmethod
615644
async def start_operation(
616645
self,
@@ -804,6 +833,35 @@ async def execute_operation(
804833
rpc_timeout: timedelta | None = None,
805834
) -> OutputT: ...
806835

836+
# Overload for temporal_operation methods
837+
@overload
838+
@abstractmethod
839+
async def execute_operation(
840+
self,
841+
operation: Callable[
842+
[
843+
NexusServiceType,
844+
temporalio.nexus.TemporalNexusStartOperationContext,
845+
temporalio.nexus.TemporalNexusClient,
846+
InputT,
847+
],
848+
Awaitable[temporalio.nexus.TemporalOperationResult[OutputT]],
849+
],
850+
arg: InputT,
851+
*,
852+
id: str,
853+
id_reuse_policy: temporalio.common.NexusOperationIDReusePolicy = temporalio.common.NexusOperationIDReusePolicy.ALLOW_DUPLICATE,
854+
id_conflict_policy: temporalio.common.NexusOperationIDConflictPolicy = temporalio.common.NexusOperationIDConflictPolicy.FAIL,
855+
schedule_to_close_timeout: timedelta | None = None,
856+
schedule_to_start_timeout: timedelta | None = None,
857+
start_to_close_timeout: timedelta | None = None,
858+
search_attributes: temporalio.common.TypedSearchAttributes | None = None,
859+
summary: str | None = None,
860+
headers: Mapping[str, str] | None = None,
861+
rpc_metadata: Mapping[str, str | bytes] = {},
862+
rpc_timeout: timedelta | None = None,
863+
) -> OutputT: ...
864+
807865
@abstractmethod
808866
async def execute_operation(
809867
self,

‎temporalio/nexus/__init__.py‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,17 @@
33
See https://github.com/temporalio/sdk-python/tree/main#nexus
44
"""
55

6-
from ._decorators import workflow_run_operation
6+
from ._decorators import (
7+
TemporalNexusOperationStartHandlerFunc,
8+
temporal_operation,
9+
workflow_run_operation,
10+
)
711
from ._operation_context import (
812
Info,
913
LoggerAdapter,
1014
NexusCallback,
15+
TemporalNexusCancelOperationContext,
16+
TemporalNexusStartOperationContext,
1117
WorkflowRunOperationContext,
1218
client,
1319
in_operation,
@@ -18,14 +24,22 @@
1824
wait_for_worker_shutdown,
1925
wait_for_worker_shutdown_sync,
2026
)
27+
from ._operation_handlers import (
28+
CancelWorkflowRunOptions,
29+
TemporalNexusOperationHandler,
30+
)
31+
from ._temporal_client import TemporalNexusClient, TemporalOperationResult
2132
from ._token import WorkflowHandle
2233

2334
__all__ = (
2435
"workflow_run_operation",
36+
"CancelWorkflowRunOptions",
2537
"Info",
2638
"LoggerAdapter",
2739
"NexusCallback",
2840
"WorkflowRunOperationContext",
41+
"TemporalNexusCancelOperationContext",
42+
"TemporalNexusStartOperationContext",
2943
"client",
3044
"in_operation",
3145
"info",
@@ -35,4 +49,9 @@
3549
"wait_for_worker_shutdown",
3650
"wait_for_worker_shutdown_sync",
3751
"WorkflowHandle",
52+
"TemporalNexusClient",
53+
"TemporalNexusOperationStartHandlerFunc",
54+
"TemporalNexusOperationHandler",
55+
"TemporalOperationResult",
56+
"temporal_operation",
3857
)

‎temporalio/nexus/_decorators.py‎

Lines changed: 135 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
from collections.abc import Awaitable, Callable
44
from typing import (
5-
TypeVar,
5+
TypeAlias,
66
overload,
77
)
88

@@ -12,27 +12,40 @@
1212
OperationHandler,
1313
StartOperationContext,
1414
)
15+
from typing_extensions import override
1516

16-
from ._operation_context import WorkflowRunOperationContext
17-
from ._operation_handlers import WorkflowRunOperationHandler
17+
from temporalio.nexus._temporal_client import (
18+
TemporalNexusClient,
19+
TemporalOperationResult,
20+
)
21+
from temporalio.types import NexusServiceType
22+
23+
from ._operation_context import (
24+
TemporalNexusStartOperationContext,
25+
WorkflowRunOperationContext,
26+
)
27+
from ._operation_handlers import (
28+
TemporalNexusOperationHandler,
29+
WorkflowRunOperationHandler,
30+
)
1831
from ._token import WorkflowHandle
1932
from ._util import (
2033
get_callable_name,
34+
get_temporal_operation_start_method_input_and_output_type_annotations,
2135
get_workflow_run_start_method_input_and_output_type_annotations,
36+
is_async_callable,
2237
set_operation_factory,
2338
)
2439

25-
ServiceHandlerT = TypeVar("ServiceHandlerT")
26-
2740

2841
@overload
2942
def workflow_run_operation(
3043
start: Callable[
31-
[ServiceHandlerT, WorkflowRunOperationContext, InputT],
44+
[NexusServiceType, WorkflowRunOperationContext, InputT],
3245
Awaitable[WorkflowHandle[OutputT]],
3346
],
3447
) -> Callable[
35-
[ServiceHandlerT, WorkflowRunOperationContext, InputT],
48+
[NexusServiceType, WorkflowRunOperationContext, InputT],
3649
Awaitable[WorkflowHandle[OutputT]],
3750
]: ...
3851

@@ -44,12 +57,12 @@ def workflow_run_operation(
4457
) -> Callable[
4558
[
4659
Callable[
47-
[ServiceHandlerT, WorkflowRunOperationContext, InputT],
60+
[NexusServiceType, WorkflowRunOperationContext, InputT],
4861
Awaitable[WorkflowHandle[OutputT]],
4962
]
5063
],
5164
Callable[
52-
[ServiceHandlerT, WorkflowRunOperationContext, InputT],
65+
[NexusServiceType, WorkflowRunOperationContext, InputT],
5366
Awaitable[WorkflowHandle[OutputT]],
5467
],
5568
]: ...
@@ -59,26 +72,26 @@ def workflow_run_operation(
5972
start: None
6073
| (
6174
Callable[
62-
[ServiceHandlerT, WorkflowRunOperationContext, InputT],
75+
[NexusServiceType, WorkflowRunOperationContext, InputT],
6376
Awaitable[WorkflowHandle[OutputT]],
6477
]
6578
) = None,
6679
*,
6780
name: str | None = None,
6881
) -> (
6982
Callable[
70-
[ServiceHandlerT, WorkflowRunOperationContext, InputT],
83+
[NexusServiceType, WorkflowRunOperationContext, InputT],
7184
Awaitable[WorkflowHandle[OutputT]],
7285
]
7386
| Callable[
7487
[
7588
Callable[
76-
[ServiceHandlerT, WorkflowRunOperationContext, InputT],
89+
[NexusServiceType, WorkflowRunOperationContext, InputT],
7790
Awaitable[WorkflowHandle[OutputT]],
7891
]
7992
],
8093
Callable[
81-
[ServiceHandlerT, WorkflowRunOperationContext, InputT],
94+
[NexusServiceType, WorkflowRunOperationContext, InputT],
8295
Awaitable[WorkflowHandle[OutputT]],
8396
],
8497
]
@@ -87,11 +100,11 @@ def workflow_run_operation(
87100

88101
def decorator(
89102
start: Callable[
90-
[ServiceHandlerT, WorkflowRunOperationContext, InputT],
103+
[NexusServiceType, WorkflowRunOperationContext, InputT],
91104
Awaitable[WorkflowHandle[OutputT]],
92105
],
93106
) -> Callable[
94-
[ServiceHandlerT, WorkflowRunOperationContext, InputT],
107+
[NexusServiceType, WorkflowRunOperationContext, InputT],
95108
Awaitable[WorkflowHandle[OutputT]],
96109
]:
97110
(
@@ -100,7 +113,7 @@ def decorator(
100113
) = get_workflow_run_start_method_input_and_output_type_annotations(start)
101114

102115
def operation_handler_factory(
103-
self: ServiceHandlerT,
116+
self: NexusServiceType,
104117
) -> OperationHandler[InputT, OutputT]:
105118
async def _start(
106119
ctx: StartOperationContext, input: InputT
@@ -130,3 +143,109 @@ async def _start(
130143
return decorator
131144

132145
return decorator(start)
146+
147+
148+
TemporalNexusOperationStartHandlerFunc: TypeAlias = Callable[
149+
[
150+
NexusServiceType,
151+
TemporalNexusStartOperationContext,
152+
TemporalNexusClient,
153+
InputT,
154+
],
155+
Awaitable[TemporalOperationResult[OutputT]],
156+
]
157+
158+
159+
@overload
160+
def temporal_operation(
161+
start: TemporalNexusOperationStartHandlerFunc[NexusServiceType, InputT, OutputT],
162+
) -> TemporalNexusOperationStartHandlerFunc[NexusServiceType, InputT, OutputT]: ...
163+
164+
165+
@overload
166+
def temporal_operation(
167+
*,
168+
name: str | None = None,
169+
) -> Callable[
170+
[TemporalNexusOperationStartHandlerFunc[NexusServiceType, InputT, OutputT]],
171+
TemporalNexusOperationStartHandlerFunc[NexusServiceType, InputT, OutputT],
172+
]: ...
173+
174+
175+
def temporal_operation(
176+
start: None
177+
| TemporalNexusOperationStartHandlerFunc[NexusServiceType, InputT, OutputT] = None,
178+
*,
179+
name: str | None = None,
180+
) -> (
181+
TemporalNexusOperationStartHandlerFunc[NexusServiceType, InputT, OutputT]
182+
| Callable[
183+
[TemporalNexusOperationStartHandlerFunc[NexusServiceType, InputT, OutputT]],
184+
TemporalNexusOperationStartHandlerFunc[NexusServiceType, InputT, OutputT],
185+
]
186+
):
187+
"""Decorator marking a method as the start method for an operation that interacts with Temporal.
188+
189+
.. warning::
190+
This API is experimental and unstable.
191+
"""
192+
193+
def decorator(
194+
start: TemporalNexusOperationStartHandlerFunc[
195+
NexusServiceType, InputT, OutputT
196+
],
197+
) -> TemporalNexusOperationStartHandlerFunc[NexusServiceType, InputT, OutputT]:
198+
if not is_async_callable(start):
199+
raise RuntimeError(
200+
f"{start} is not an `async def` method. "
201+
"@temporal_operation must decorate an `async def` start method."
202+
)
203+
(
204+
input_type,
205+
output_type,
206+
) = get_temporal_operation_start_method_input_and_output_type_annotations(start)
207+
208+
def operation_handler_factory(
209+
self: NexusServiceType,
210+
) -> OperationHandler[InputT, OutputT]:
211+
async def _start(
212+
ctx: TemporalNexusStartOperationContext,
213+
client: TemporalNexusClient,
214+
input: InputT,
215+
) -> TemporalOperationResult[OutputT]:
216+
return await start(
217+
self,
218+
ctx,
219+
client,
220+
input,
221+
)
222+
223+
class _TemporalNexusOperationHandler(TemporalNexusOperationHandler):
224+
@override
225+
async def start_operation(
226+
self,
227+
ctx: TemporalNexusStartOperationContext,
228+
client: TemporalNexusClient,
229+
input: InputT,
230+
) -> TemporalOperationResult[OutputT]:
231+
return await _start(ctx, client, input)
232+
233+
_TemporalNexusOperationHandler.start_operation.__doc__ = start.__doc__
234+
return _TemporalNexusOperationHandler()
235+
236+
method_name = get_callable_name(start)
237+
op = nexusrpc.Operation(
238+
name=name or method_name,
239+
input_type=input_type,
240+
output_type=output_type,
241+
)
242+
op.method_name = method_name
243+
nexusrpc.set_operation(operation_handler_factory, op)
244+
245+
set_operation_factory(start, operation_handler_factory)
246+
return start
247+
248+
if start is None:
249+
return decorator
250+
251+
return decorator(start)

0 commit comments

Comments
 (0)