1
1
import json
2
- from typing import Any , Coroutine , Generic , List , Optional , TypeVar , Union , overload
2
+ from typing import (
3
+ Any ,
4
+ Coroutine ,
5
+ Generic ,
6
+ List ,
7
+ Literal ,
8
+ Optional ,
9
+ TypeVar ,
10
+ Union ,
11
+ overload ,
12
+ )
3
13
4
14
import pydantic
5
15
from taskiq import AsyncBroker , AsyncTaskiqTask
6
16
from taskiq .decor import AsyncTaskiqDecoratedTask
7
17
from taskiq .kicker import AsyncKicker
8
18
from typing_extensions import ParamSpec
9
19
10
- from taskiq_pipelines .constants import CURRENT_STEP , PIPELINE_DATA
20
+ from taskiq_pipelines .constants import CURRENT_STEP , EMPTY_PARAM_NAME , PIPELINE_DATA
11
21
from taskiq_pipelines .steps import FilterStep , MapperStep , SequentialStep , parse_step
12
22
13
23
_ReturnType = TypeVar ("_ReturnType" )
@@ -58,7 +68,7 @@ def call_next(
58
68
AsyncKicker [Any , Coroutine [Any , Any , _T2 ]],
59
69
AsyncTaskiqDecoratedTask [Any , Coroutine [Any , Any , _T2 ]],
60
70
],
61
- param_name : Optional [str ] = None ,
71
+ param_name : Union [ Optional [str ], Literal [ - 1 ] ] = None ,
62
72
** additional_kwargs : Any ,
63
73
) -> "Pipeline[_FuncParams, _T2]" :
64
74
...
@@ -70,7 +80,7 @@ def call_next(
70
80
AsyncKicker [Any , _T2 ],
71
81
AsyncTaskiqDecoratedTask [Any , _T2 ],
72
82
],
73
- param_name : Optional [str ] = None ,
83
+ param_name : Union [ Optional [str ], Literal [ - 1 ] ] = None ,
74
84
** additional_kwargs : Any ,
75
85
) -> "Pipeline[_FuncParams, _T2]" :
76
86
...
@@ -81,7 +91,7 @@ def call_next(
81
91
AsyncKicker [Any , Any ],
82
92
AsyncTaskiqDecoratedTask [Any , Any ],
83
93
],
84
- param_name : Optional [str ] = None ,
94
+ param_name : Union [ Optional [str ], Literal [ - 1 ] ] = None ,
85
95
** additional_kwargs : Any ,
86
96
) -> Any :
87
97
"""
@@ -94,13 +104,14 @@ def call_next(
94
104
if param_name is specified.
95
105
96
106
:param task: task to execute.
97
- :param param_name: kwarg param name, defaults to None
107
+ :param param_name: kwarg param name, defaults to None.
108
+ If set to -1 (EMPTY_PARAM_NAME), result is not passed.
98
109
:param additional_kwargs: additional kwargs to task.
99
110
:return: updated pipeline.
100
111
"""
101
112
self .steps .append (
102
113
DumpedStep (
103
- step_type = SequentialStep .step_name ,
114
+ step_type = SequentialStep ._step_name , # noqa: WPS437
104
115
step_data = SequentialStep .from_task (
105
116
task = task ,
106
117
param_name = param_name ,
@@ -111,6 +122,62 @@ def call_next(
111
122
)
112
123
return self
113
124
125
+ @overload
126
+ def call_after (
127
+ self : "Pipeline[_FuncParams, _ReturnType]" ,
128
+ task : Union [
129
+ AsyncKicker [Any , Coroutine [Any , Any , _T2 ]],
130
+ AsyncTaskiqDecoratedTask [Any , Coroutine [Any , Any , _T2 ]],
131
+ ],
132
+ ** additional_kwargs : Any ,
133
+ ) -> "Pipeline[_FuncParams, _T2]" :
134
+ ...
135
+
136
+ @overload
137
+ def call_after (
138
+ self : "Pipeline[_FuncParams, _ReturnType]" ,
139
+ task : Union [
140
+ AsyncKicker [Any , _T2 ],
141
+ AsyncTaskiqDecoratedTask [Any , _T2 ],
142
+ ],
143
+ ** additional_kwargs : Any ,
144
+ ) -> "Pipeline[_FuncParams, _T2]" :
145
+ ...
146
+
147
+ def call_after (
148
+ self ,
149
+ task : Union [
150
+ AsyncKicker [Any , Any ],
151
+ AsyncTaskiqDecoratedTask [Any , Any ],
152
+ ],
153
+ ** additional_kwargs : Any ,
154
+ ) -> Any :
155
+ """
156
+ Adds sequential step.
157
+
158
+ This task will be executed right after
159
+ the previous and result of the previous task
160
+ is not passed to the next task.
161
+
162
+ This is equivalent to call_next(task, param_name=-1).
163
+
164
+ :param task: task to execute.
165
+ :param additional_kwargs: additional kwargs to task.
166
+ :return: updated pipeline.
167
+ """
168
+ self .steps .append (
169
+ DumpedStep (
170
+ step_type = SequentialStep ._step_name , # noqa: WPS437
171
+ step_data = SequentialStep .from_task (
172
+ task = task ,
173
+ param_name = EMPTY_PARAM_NAME ,
174
+ ** additional_kwargs ,
175
+ ).dumps (),
176
+ task_id = "" ,
177
+ ),
178
+ )
179
+ return self
180
+
114
181
@overload
115
182
def map (
116
183
self : "Pipeline[_FuncParams, _ReturnType]" ,
@@ -169,7 +236,7 @@ def map(
169
236
"""
170
237
self .steps .append (
171
238
DumpedStep (
172
- step_type = MapperStep .step_name ,
239
+ step_type = MapperStep ._step_name , # noqa: WPS437
173
240
step_data = MapperStep .from_task (
174
241
task = task ,
175
242
param_name = param_name ,
@@ -241,7 +308,7 @@ def filter(
241
308
"""
242
309
self .steps .append (
243
310
DumpedStep (
244
- step_type = FilterStep .step_name ,
311
+ step_type = FilterStep ._step_name , # noqa: WPS437
245
312
step_data = FilterStep .from_task (
246
313
task = task ,
247
314
param_name = param_name ,
@@ -261,7 +328,7 @@ def dumps(self) -> str:
261
328
:returns: serialized pipeline.
262
329
"""
263
330
return json .dumps (
264
- [step .dict () for step in self .steps ],
331
+ [step .model_dump () for step in self .steps ],
265
332
)
266
333
267
334
@classmethod
@@ -277,7 +344,7 @@ def loads(cls, broker: AsyncBroker, pipe_data: str) -> "Pipeline[Any, Any]":
277
344
:return: new
278
345
"""
279
346
pipe : "Pipeline[Any, Any]" = Pipeline (broker )
280
- pipe .steps = pydantic .parse_raw_as (List [DumpedStep ], pipe_data )
347
+ pipe .steps = pydantic .TypeAdapter (List [DumpedStep ]). validate_json ( pipe_data )
281
348
return pipe
282
349
283
350
async def kiq (
0 commit comments