Skip to content

Commit eb5daff

Browse files
dandavisondandavison
committed
Code golf
1 parent 880b487 commit eb5daff

File tree

1 file changed

+53
-52
lines changed

1 file changed

+53
-52
lines changed

‎temporalio/worker/_command_aware_visitor.py‎

Lines changed: 53 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
"""Visitor that sets command context during payload traversal."""
22

33
import contextvars
4+
import sys
45
from contextlib import contextmanager
56
from dataclasses import dataclass
67
from typing import Any, Iterator, Optional, Type
@@ -51,103 +52,103 @@ class CommandAwarePayloadVisitor(PayloadVisitor):
5152
async def _visit_coresdk_workflow_commands_ScheduleActivity(
5253
self, fs: VisitorFunctions, o: ScheduleActivity
5354
) -> None:
54-
with current_command(CommandType.COMMAND_TYPE_SCHEDULE_ACTIVITY_TASK, o.seq):
55-
await super()._visit_coresdk_workflow_commands_ScheduleActivity(fs, o)
55+
await self._visit_with_context(
56+
CommandType.COMMAND_TYPE_SCHEDULE_ACTIVITY_TASK, o.seq, fs, o
57+
)
5658

5759
async def _visit_coresdk_workflow_commands_ScheduleLocalActivity(
5860
self, fs: VisitorFunctions, o: ScheduleLocalActivity
5961
) -> None:
60-
with current_command(CommandType.COMMAND_TYPE_SCHEDULE_ACTIVITY_TASK, o.seq):
61-
await super()._visit_coresdk_workflow_commands_ScheduleLocalActivity(fs, o)
62+
await self._visit_with_context(
63+
CommandType.COMMAND_TYPE_SCHEDULE_ACTIVITY_TASK, o.seq, fs, o
64+
)
6265

6366
async def _visit_coresdk_workflow_commands_StartChildWorkflowExecution(
6467
self, fs: VisitorFunctions, o: StartChildWorkflowExecution
6568
) -> None:
66-
with current_command(
67-
CommandType.COMMAND_TYPE_START_CHILD_WORKFLOW_EXECUTION, o.seq
68-
):
69-
await super()._visit_coresdk_workflow_commands_StartChildWorkflowExecution(
70-
fs, o
71-
)
69+
await self._visit_with_context(
70+
CommandType.COMMAND_TYPE_START_CHILD_WORKFLOW_EXECUTION, o.seq, fs, o
71+
)
7272

7373
async def _visit_coresdk_workflow_commands_SignalExternalWorkflowExecution(
7474
self, fs: VisitorFunctions, o: SignalExternalWorkflowExecution
7575
) -> None:
76-
with current_command(
77-
CommandType.COMMAND_TYPE_SIGNAL_EXTERNAL_WORKFLOW_EXECUTION, o.seq
78-
):
79-
await super()._visit_coresdk_workflow_commands_SignalExternalWorkflowExecution(
80-
fs, o
81-
)
76+
await self._visit_with_context(
77+
CommandType.COMMAND_TYPE_SIGNAL_EXTERNAL_WORKFLOW_EXECUTION, o.seq, fs, o
78+
)
8279

8380
async def _visit_coresdk_workflow_commands_ScheduleNexusOperation(
8481
self, fs: VisitorFunctions, o: ScheduleNexusOperation
8582
) -> None:
86-
with current_command(CommandType.COMMAND_TYPE_SCHEDULE_NEXUS_OPERATION, o.seq):
87-
await super()._visit_coresdk_workflow_commands_ScheduleNexusOperation(fs, o)
83+
await self._visit_with_context(
84+
CommandType.COMMAND_TYPE_SCHEDULE_NEXUS_OPERATION, o.seq, fs, o
85+
)
8886

8987
# Workflow activation jobs with payloads
9088
async def _visit_coresdk_workflow_activation_ResolveActivity(
9189
self, fs: VisitorFunctions, o: ResolveActivity
9290
) -> None:
93-
with current_command(CommandType.COMMAND_TYPE_SCHEDULE_ACTIVITY_TASK, o.seq):
94-
await super()._visit_coresdk_workflow_activation_ResolveActivity(fs, o)
91+
await self._visit_with_context(
92+
CommandType.COMMAND_TYPE_SCHEDULE_ACTIVITY_TASK, o.seq, fs, o
93+
)
9594

9695
async def _visit_coresdk_workflow_activation_ResolveChildWorkflowExecutionStart(
9796
self, fs: VisitorFunctions, o: ResolveChildWorkflowExecutionStart
9897
) -> None:
99-
with current_command(
100-
CommandType.COMMAND_TYPE_START_CHILD_WORKFLOW_EXECUTION, o.seq
101-
):
102-
await super()._visit_coresdk_workflow_activation_ResolveChildWorkflowExecutionStart(
103-
fs, o
104-
)
98+
await self._visit_with_context(
99+
CommandType.COMMAND_TYPE_START_CHILD_WORKFLOW_EXECUTION, o.seq, fs, o
100+
)
105101

106102
async def _visit_coresdk_workflow_activation_ResolveChildWorkflowExecution(
107103
self, fs: VisitorFunctions, o: ResolveChildWorkflowExecution
108104
) -> None:
109-
with current_command(
110-
CommandType.COMMAND_TYPE_START_CHILD_WORKFLOW_EXECUTION, o.seq
111-
):
112-
await super()._visit_coresdk_workflow_activation_ResolveChildWorkflowExecution(
113-
fs, o
114-
)
105+
await self._visit_with_context(
106+
CommandType.COMMAND_TYPE_START_CHILD_WORKFLOW_EXECUTION, o.seq, fs, o
107+
)
115108

116109
async def _visit_coresdk_workflow_activation_ResolveSignalExternalWorkflow(
117110
self, fs: VisitorFunctions, o: ResolveSignalExternalWorkflow
118111
) -> None:
119-
with current_command(
120-
CommandType.COMMAND_TYPE_SIGNAL_EXTERNAL_WORKFLOW_EXECUTION, o.seq
121-
):
122-
await super()._visit_coresdk_workflow_activation_ResolveSignalExternalWorkflow(
123-
fs, o
124-
)
112+
await self._visit_with_context(
113+
CommandType.COMMAND_TYPE_SIGNAL_EXTERNAL_WORKFLOW_EXECUTION, o.seq, fs, o
114+
)
125115

126116
async def _visit_coresdk_workflow_activation_ResolveRequestCancelExternalWorkflow(
127117
self, fs: VisitorFunctions, o: ResolveRequestCancelExternalWorkflow
128118
) -> None:
129-
with current_command(
130-
CommandType.COMMAND_TYPE_REQUEST_CANCEL_EXTERNAL_WORKFLOW_EXECUTION, o.seq
131-
):
132-
await super()._visit_coresdk_workflow_activation_ResolveRequestCancelExternalWorkflow(
133-
fs, o
134-
)
119+
await self._visit_with_context(
120+
CommandType.COMMAND_TYPE_REQUEST_CANCEL_EXTERNAL_WORKFLOW_EXECUTION,
121+
o.seq,
122+
fs,
123+
o,
124+
)
135125

136126
async def _visit_coresdk_workflow_activation_ResolveNexusOperationStart(
137127
self, fs: VisitorFunctions, o: ResolveNexusOperationStart
138128
) -> None:
139-
with current_command(CommandType.COMMAND_TYPE_SCHEDULE_NEXUS_OPERATION, o.seq):
140-
await super()._visit_coresdk_workflow_activation_ResolveNexusOperationStart(
141-
fs, o
142-
)
129+
await self._visit_with_context(
130+
CommandType.COMMAND_TYPE_SCHEDULE_NEXUS_OPERATION, o.seq, fs, o
131+
)
143132

144133
async def _visit_coresdk_workflow_activation_ResolveNexusOperation(
145134
self, fs: VisitorFunctions, o: ResolveNexusOperation
146135
) -> None:
147-
with current_command(CommandType.COMMAND_TYPE_SCHEDULE_NEXUS_OPERATION, o.seq):
148-
await super()._visit_coresdk_workflow_activation_ResolveNexusOperation(
149-
fs, o
150-
)
136+
await self._visit_with_context(
137+
CommandType.COMMAND_TYPE_SCHEDULE_NEXUS_OPERATION, o.seq, fs, o
138+
)
139+
140+
async def _visit_with_context(
141+
self,
142+
command_type: CommandType.ValueType,
143+
seq: int,
144+
fs: VisitorFunctions,
145+
o: Any,
146+
) -> None:
147+
"""Helper to call parent method with command context."""
148+
method_name = sys._getframe(1).f_code.co_name
149+
parent_method = getattr(super(), method_name)
150+
with current_command(command_type, seq):
151+
await parent_method(fs, o)
151152

152153

153154
def _get_workflow_command_protos_with_seq() -> Iterator[Type[Any]]:

0 commit comments

Comments
 (0)