Initial groups support.

Author: s3rius

README.md

@@ -162,3 +162,14 @@ If called tasks returned `True` for some element, this element will be added in
 
 After the execution you'll get a list with filtered results.
 You can add filters by calling `.filter` method of the pipeline.
+
+
+### Group steps
+
+This step groups together multiple tasks and sends them after the previous steps.
+
+To create a group you need to use `Group` class from `taskiq_pipelienes` like this:
+
+```
+
+```

taskiq_pipelines/abc.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
 from typing import Any, Dict, Type
 
-from taskiq import AsyncBroker, TaskiqResult
+from taskiq import AsyncBroker, AsyncTaskiqTask, TaskiqResult
 from typing_extensions import ClassVar
 
 
@@ -26,9 +26,9 @@ async def act(
         step_number: int,
         parent_task_id: str,
         task_id: str,
-        pipe_data: str,
+        pipe_data: bytes,
         result: "TaskiqResult[Any]",
-    ) -> None:
+    ) -> AsyncTaskiqTask[Any]:
         """
         Perform pipeline action.
 

taskiq_pipelines/pipeliner.py

@@ -13,7 +13,7 @@
 )
 
 import pydantic
-from taskiq import AsyncBroker, AsyncTaskiqTask
+from taskiq import AsyncBroker, AsyncTaskiqTask, TaskiqResult
 from taskiq.decor import AsyncTaskiqDecoratedTask
 from taskiq.kicker import AsyncKicker
 from typing_extensions import ParamSpec
@@ -52,20 +52,44 @@ class Pipeline(Generic[_FuncParams, _ReturnType]):
     but it's nice to have.
     """
 
+    @overload
+    def __init__(
+        self: "Pipeline[[], _ReturnType]",
+        broker: AsyncBroker,
+        task: Optional[Group[_ReturnType]] = None,
+    ) -> None: ...
+
+    @overload
+    def __init__(
+        self,
+        broker: AsyncBroker,
+        task: Optional[
+            Union[
+                AsyncKicker[_FuncParams, _ReturnType],
+                AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType],
+            ]
+        ] = None,
+    ) -> None: ...
+
     def __init__(
         self,
         broker: AsyncBroker,
         task: Optional[
             Union[
                 AsyncKicker[_FuncParams, _ReturnType],
                 AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType],
+                Group[_ReturnType],
             ]
         ] = None,
     ) -> None:
         self.broker = broker
         self.steps: "List[DumpedStep]" = []
-        if task:
-            self.call_next(task)
+        if not task:
+            return
+        if isinstance(task, Group):
+            self.group(task)
+            return
+        self.call_next(task)
 
     @overload
     def call_next(
@@ -379,6 +403,48 @@ def loadb(cls, broker: AsyncBroker, pipe_data: bytes) -> "Pipeline[Any, Any]":
         pipe.steps = DumpedSteps.model_validate(data)  # type: ignore[assignment]
         return pipe
 
+    async def _kick_sequential(
+        self,
+        step: SequentialStep,
+        task_id: str,
+        *args: Any,
+        **kwargs: Any,
+    ) -> AsyncTaskiqTask[_ReturnType]:
+        kicker = (
+            AsyncKicker(
+                step.task_name,
+                broker=self.broker,
+                labels=step.labels,
+            )
+            .with_task_id(task_id)
+            .with_labels(
+                **{CURRENT_STEP: 0, PIPELINE_DATA: self.dumpb()},  # type: ignore
+            )
+        )
+        return await kicker.kiq(*args, **kwargs)
+
+    async def _kick_group(
+        self,
+        group: GroupStep,
+        task_id: str,
+    ) -> AsyncTaskiqTask[Any]:
+        await group.act(
+            broker=self.broker,
+            task_id=task_id,
+            step_number=0,
+            parent_task_id="",
+            pipe_data=self.dumpb(),
+            result=TaskiqResult(
+                is_err=False,
+                return_value=None,
+                execution_time=0.0,
+            ),
+        )
+        return AsyncTaskiqTask(
+            task_id=task_id,
+            result_backend=self.broker.result_backend,
+        )
+
     async def kiq(
         self,
         *args: _FuncParams.args,
@@ -405,20 +471,18 @@ async def kiq(
         self._update_task_ids()
         step = self.steps[0]
         parsed_step = parse_step(step.step_type, step.step_data)
-        if not isinstance(parsed_step, SequentialStep):
-            raise ValueError("First step must be sequential.")
-        kicker = (
-            AsyncKicker(
-                parsed_step.task_name,
-                broker=self.broker,
-                labels=parsed_step.labels,
+        if isinstance(parsed_step, SequentialStep):
+            taskiq_task = await self._kick_sequential(
+                parsed_step,
+                step.task_id,
+                *args,
+                **kwargs,
             )
-            .with_task_id(step.task_id)
-            .with_labels(
-                **{CURRENT_STEP: 0, PIPELINE_DATA: self.dumpb()},  # type: ignore
-            )
-        )
-        taskiq_task = await kicker.kiq(*args, **kwargs)
+        elif isinstance(parsed_step, GroupStep):
+            taskiq_task = await self._kick_group(parsed_step, step.task_id)
+        else:
+            raise ValueError("First step must be sequential or a group.")
+
         taskiq_task.task_id = self.steps[-1].task_id
         return taskiq_task
 

taskiq_pipelines/steps/filter.py

@@ -2,7 +2,7 @@
 from typing import Any, Dict, Iterable, List, Optional, Union
 
 import pydantic
-from taskiq import AsyncBroker, Context, TaskiqDepends, TaskiqResult
+from taskiq import AsyncBroker, AsyncTaskiqTask, Context, TaskiqDepends, TaskiqResult
 from taskiq.brokers.shared_broker import async_shared_broker
 from taskiq.decor import AsyncTaskiqDecoratedTask
 from taskiq.kicker import AsyncKicker
@@ -88,9 +88,9 @@ async def act(
         step_number: int,
         parent_task_id: str,
         task_id: str,
-        pipe_data: str,
+        pipe_data: bytes,
         result: "TaskiqResult[Any]",
-    ) -> None:
+    ) -> AsyncTaskiqTask[Any]:
         """
         Run filter action.
 
@@ -121,7 +121,7 @@ async def act(
             else:
                 task = await kicker.kiq(item, **self.additional_kwargs)
             sub_task_ids.append(task.task_id)
-        await (
+        return await (
             filter_tasks.kicker()
             .with_task_id(task_id)
             .with_broker(

taskiq_pipelines/steps/group.py

@@ -3,6 +3,7 @@
 import pydantic
 from taskiq import (
     AsyncBroker,
+    AsyncTaskiqTask,
     Context,
     TaskiqDepends,
     TaskiqMessage,
@@ -82,9 +83,9 @@ async def act(
         step_number: int,
         parent_task_id: str,
         task_id: str,
-        pipe_data: str,
+        pipe_data: bytes,
         result: "TaskiqResult[Any]",
-    ) -> None:
+    ) -> AsyncTaskiqTask[Any]:
         """
         Execute group action.
 
@@ -100,7 +101,7 @@ async def act(
             ids.append(subtask_id)
             await broker.kick(broker.formatter.dumps(task.to_message(subtask_id)))
 
-        await (
+        return await (
             wait_group_tasks.kicker()
             .with_broker(broker)
             .with_task_id(task_id)

taskiq_pipelines/steps/mapper.py

@@ -6,6 +6,7 @@
 from taskiq import (
     AsyncBroker,
     AsyncTaskiqDecoratedTask,
+    AsyncTaskiqTask,
     Context,
     TaskiqDepends,
     TaskiqResult,
@@ -60,8 +61,8 @@ async def wait_tasks(  # noqa: C901
     results: List[Any] = []
     for task_id in ordered_ids:
         result = await context.broker.result_backend.get_result(task_id)
-        logger.warning("Found error: %s", result.error)
         if result.is_err:
+            logger.warning("Found error: %s", result.error)
             if skip_errors:
                 continue
             if none_if_errors:
@@ -92,9 +93,9 @@ async def act(
         step_number: int,
         parent_task_id: str,
         task_id: str,
-        pipe_data: str,
+        pipe_data: bytes,
         result: "TaskiqResult[Any]",
-    ) -> None:
+    ) -> AsyncTaskiqTask[Any]:
         """
         Runs mapping.
 
@@ -132,7 +133,7 @@ async def act(
                 task = await kicker.kiq(item, **self.additional_kwargs)
             sub_task_ids.append(task.task_id)
 
-        await (
+        return await (
             wait_tasks.kicker()
             .with_task_id(task_id)
             .with_broker(

taskiq_pipelines/steps/sequential.py

@@ -1,7 +1,7 @@
 from typing import Any, Dict, Optional, Union
 
 import pydantic
-from taskiq import AsyncBroker, AsyncTaskiqDecoratedTask, TaskiqResult
+from taskiq import AsyncBroker, AsyncTaskiqDecoratedTask, AsyncTaskiqTask, TaskiqResult
 from taskiq.kicker import AsyncKicker
 
 from taskiq_pipelines.abc import AbstractStep
@@ -30,9 +30,9 @@ async def act(
         step_number: int,
         parent_task_id: str,
         task_id: str,
-        pipe_data: str,
+        pipe_data: bytes,
         result: "TaskiqResult[Any]",
-    ) -> None:
+    ) -> AsyncTaskiqTask[Any]:
         """
         Runs next task.
 
@@ -64,11 +64,10 @@ async def act(
         )
         if isinstance(self.param_name, str):
             self.additional_kwargs[self.param_name] = result.return_value
-            await kicker.kiq(**self.additional_kwargs)
-        elif self.param_name == EMPTY_PARAM_NAME:
-            await kicker.kiq(**self.additional_kwargs)
-        else:
-            await kicker.kiq(result.return_value, **self.additional_kwargs)
+            return await kicker.kiq(**self.additional_kwargs)
+        if self.param_name == EMPTY_PARAM_NAME:
+            return await kicker.kiq(**self.additional_kwargs)
+        return await kicker.kiq(result.return_value, **self.additional_kwargs)
 
     @classmethod
     def from_task(