Added group with args.

Author: s3rius

taskiq_pipelines/pipeliner.py

@@ -21,7 +21,7 @@
 from taskiq_pipelines.constants import CURRENT_STEP, EMPTY_PARAM_NAME, PIPELINE_DATA
 from taskiq_pipelines.steps import FilterStep, MapperStep, SequentialStep, parse_step
 from taskiq_pipelines.steps.group import GroupStep
-from taskiq_pipelines.task_group import Group
+from taskiq_pipelines.task_group import Group, GroupWithArgs
 
 _ReturnType = TypeVar("_ReturnType")
 _FuncParams = ParamSpec("_FuncParams")
@@ -54,9 +54,9 @@ class Pipeline(Generic[_FuncParams, _ReturnType]):
 
     @overload
     def __init__(
-        self: "Pipeline[[], _ReturnType]",
+        self: "Pipeline[[], Any]",
         broker: AsyncBroker,
-        task: Optional[Group[Any, _ReturnType]] = None,
+        task: None = None,
     ) -> None: ...
 
     @overload
@@ -351,10 +351,22 @@ def filter(
         )
         return self
 
+    @overload
+    def group(
+        self: "Pipeline[_FuncParams, _ReturnType]",
+        group: GroupWithArgs[Any, _T2],
+    ) -> "Pipeline[_FuncParams, _T2]": ...
+
+    @overload
     def group(
         self: "Pipeline[_FuncParams, _ReturnType]",
         group: Group[Any, _T2],
-    ) -> "Pipeline[_FuncParams, _T2]":
+    ) -> "Pipeline[_FuncParams, _T2]": ...
+
+    def group(
+        self: "Pipeline[_FuncParams, _ReturnType]",
+        group: Group[Any, Any] | GroupWithArgs[Any, Any],
+    ) -> "Pipeline[_FuncParams, Any]":
         """
         Add group task execution step.
 

taskiq_pipelines/steps/group.py

@@ -41,6 +41,7 @@ class GroupStepItem(pydantic.BaseModel):
     labels: Dict[str, Any]
     labels_types: Optional[Dict[str, int]] = None
     args: List[Any]
+    param_name: Optional[str]
     kwargs: Dict[str, Any]
 
     def from_message(self, message: TaskiqMessage) -> None:
@@ -54,19 +55,31 @@ def from_message(self, message: TaskiqMessage) -> None:
         self.args = message.args
         self.kwargs = message.kwargs
 
-    def to_message(self, task_id: str) -> TaskiqMessage:
+    def to_message(
+        self,
+        task_id: str,
+        result: Optional[TaskiqResult[Any]] = None,
+    ) -> TaskiqMessage:
         """
         Convert this item to message.
 
         :return: message
         """
+        args = self.args
+        kwargs = self.kwargs
+        if result:
+            if self.param_name:
+                kwargs[self.param_name] = result.return_value
+            else:
+                args = [result.return_value, *args]
+
         return TaskiqMessage(
             task_id=task_id,
             task_name=self.task_name,
             labels=self.labels,
             labels_types=self.labels_types,
-            args=self.args,
-            kwargs=self.kwargs,
+            args=args,
+            kwargs=kwargs,
         )
 
 
@@ -76,6 +89,7 @@ class GroupStep(pydantic.BaseModel, AbstractStep, step_name="group"):
     tasks: list[GroupStepItem]
     skip_errors: bool
     check_interval: float
+    pass_args: bool = False
 
     async def act(
         self,
@@ -99,7 +113,13 @@ async def act(
         for task in self.tasks:
             subtask_id = broker.id_generator()
             ids.append(subtask_id)
-            await broker.kick(broker.formatter.dumps(task.to_message(subtask_id)))
+
+            if self.pass_args:
+                message = task.to_message(subtask_id, result)
+            else:
+                message = task.to_message(subtask_id, None)
+
+            await broker.kick(broker.formatter.dumps(message))
 
         return await (
             wait_group_tasks.kicker()

taskiq_pipelines/task_group.py

@@ -113,6 +113,7 @@ def add(
                 labels_types=message.labels_types,
                 args=message.args,
                 kwargs=message.kwargs,
+                param_name=None,
             ),
         )
         return self
@@ -123,4 +124,123 @@ def to_step(self) -> GroupStep:
             tasks=list(self.tasks),
             skip_errors=self.skip_errors,
             check_interval=self.check_interval,
+            pass_args=False,
+        )
+
+
+class GroupWithArgs(Generic[_S, _T]):
+    """
+    Group of tasks with args.
+
+    This class gathers multiple tasks together.
+    They will run in parallel. The only difference
+    with Group is that it passes the result
+    of the previous task as the first argument
+    to all of the tasks in the group.
+
+    :param skip_errors: If True, errors in one task will not affect others.
+    """
+
+    # If skip_errors is set to True,
+    # We update first generic type to Literal[True].
+    @overload
+    def __init__(
+        self: "GroupWithArgs[Literal[True], Tuple[()]]",
+        skip_errors: Literal[True],
+        check_interval: float = 0.1,
+    ) -> None: ...
+
+    @overload
+    def __init__(
+        self: "GroupWithArgs[Literal[False], Tuple[()]]",
+        skip_errors: bool = False,
+        check_interval: float = 0.1,
+    ) -> None: ...
+
+    def __init__(
+        self: "GroupWithArgs[Any, Tuple[()]]",
+        skip_errors: bool = False,
+        check_interval: float = 0.1,
+    ) -> None:
+        self.tasks: Tuple[GroupStepItem, ...] = ()
+        self.skip_errors = skip_errors
+        self.check_interval = check_interval
+
+    @overload
+    def add(
+        self: "GroupWithArgs[Literal[True], Tuple[Unpack[_Tups]]]",
+        task: Union[
+            AsyncKicker[_Params, Coroutine[Any, Any, _TVal]],
+            AsyncKicker[_Params, "CoroutineType[Any, Any, _TVal]"],
+            AsyncTaskiqDecoratedTask[_Params, Coroutine[Any, Any, _TVal]],
+            AsyncTaskiqDecoratedTask[_Params, "CoroutineType[Any, Any, _TVal]"],
+        ],
+        param_name: Optional[str] = None,
+        **additional_kwargs: Any,
+    ) -> "GroupWithArgs[_S, Tuple[Unpack[_Tups], Optional[_TVal]]]": ...
+
+    @overload
+    def add(
+        self: "GroupWithArgs[Literal[False], Tuple[Unpack[_Tups]]]",
+        task: Union[
+            AsyncKicker[_Params, Coroutine[Any, Any, _TVal]],
+            AsyncKicker[_Params, "CoroutineType[Any, Any, _TVal]"],
+            AsyncTaskiqDecoratedTask[_Params, Coroutine[Any, Any, _TVal]],
+            AsyncTaskiqDecoratedTask[_Params, "CoroutineType[Any, Any, _TVal]"],
+        ],
+        param_name: Optional[str] = None,
+        **additional_kwargs: Any,
+    ) -> "GroupWithArgs[_S, Tuple[Unpack[_Tups], _TVal]]": ...
+
+    @overload
+    def add(
+        self: "GroupWithArgs[Literal[True], Tuple[Unpack[_Tups]]]",
+        task: Union[
+            AsyncKicker[_Params, _TVal],
+            AsyncTaskiqDecoratedTask[_Params, _TVal],
+        ],
+        param_name: Optional[str] = None,
+        **additional_kwargs: Any,
+    ) -> "GroupWithArgs[_S, Tuple[Unpack[_Tups], Optional[_TVal]]]": ...
+
+    @overload
+    def add(
+        self: "GroupWithArgs[Literal[False], Tuple[Unpack[_Tups]]]",
+        task: Union[
+            AsyncKicker[_Params, _TVal],
+            AsyncTaskiqDecoratedTask[_Params, _TVal],
+        ],
+        param_name: Optional[str] = None,
+        **additional_kwargs: Any,
+    ) -> "GroupWithArgs[_S, Tuple[Unpack[_Tups], _TVal]]": ...
+
+    def add(
+        self: "GroupWithArgs[Any, Any]",
+        task: Union[AsyncKicker[_Params, Any], AsyncTaskiqDecoratedTask[_Params, Any]],
+        param_name: Optional[str] = None,
+        **additional_kwargs: Any,
+    ) -> "Any":
+        """Add task to a group."""
+        kicker = task.kicker() if isinstance(task, AsyncTaskiqDecoratedTask) else task
+        message = kicker._prepare_message(**additional_kwargs)
+        self.tasks = (
+            *self.tasks,
+            GroupStepItem(
+                task_name=message.task_name,
+                labels=message.labels,
+                labels_types=message.labels_types,
+                args=message.args,
+                kwargs=message.kwargs,
+                param_name=param_name,
+            ),
+        )
+        return self
+
+    def to_step(self) -> GroupStep:
+        """Convert group definition to a step."""
+        return GroupStep(
+            tasks=list(self.tasks),
+            skip_errors=self.skip_errors,
+            check_interval=self.check_interval,
+            pass_args=True,
         )