Merge pull request #7 from taskiq-python/feature/post-save

Middleware now react on post-save hook.

Author: s3rius

poetry.lock

@@ -6,7 +6,7 @@ authors = ["Pavel Kirilin <[email protected]>"]
 
 [tool.poetry.dependencies]
 python = "^3.7"
-taskiq = "^0"
+taskiq = ">=0.0.8, <1"
 typing-extensions = "^4.3.0"
 pydantic = "^1.6.2"
 
@@ -22,7 +22,7 @@ isort = "^5.10.1"
 yesqa = "^1.4.0"
 wemake-python-styleguide = "^0.16.1"
 mypy = "^0.971"
-pytest-xdist = {version = "^2.5.0", extras = ["psutil"]}
+pytest-xdist = { version = "^2.5.0", extras = ["psutil"] }
 
 [tool.mypy]
 strict = true

pyproject.toml

@@ -1 +1,11 @@
 """Pipelines for taskiq tasks."""
+from taskiq_pipelines.exceptions import AbortPipeline, PipelineError
+from taskiq_pipelines.middleware import PipelineMiddleware
+from taskiq_pipelines.pipeliner import Pipeline
+
+__all__ = [
+    "Pipeline",
+    "PipelineError",
+    "AbortPipeline",
+    "PipelineMiddleware",
+]

taskiq_pipelines/__init__.py

@@ -15,7 +15,7 @@
 class PipelineMiddleware(TaskiqMiddleware):
     """Pipeline middleware."""
 
-    async def post_execute(  # noqa: C901, WPS212
+    async def post_save(  # noqa: C901, WPS212
         self,
         message: "TaskiqMessage",
         result: "TaskiqResult[Any]",

taskiq_pipelines/middleware.py

@@ -50,7 +50,8 @@ async def filter_tasks(  # noqa: C901, WPS210, WPS231
                     tasks_set.remove(task_id)
                 except LookupError:
                     continue
-        await asyncio.sleep(check_interval)
+        if tasks_set:
+            await asyncio.sleep(check_interval)
 
     results = await context.broker.result_backend.get_result(parent_task_id)
     filtered_results = []
@@ -135,7 +136,6 @@ async def act(
             else:
                 task = await kicker.kiq(item, **self.additional_kwargs)
             sub_task_ids.append(task.task_id)
-
         await filter_tasks.kicker().with_task_id(task_id).with_broker(
             broker,
         ).with_labels(

taskiq_pipelines/steps/filter.py

@@ -50,7 +50,8 @@ async def wait_tasks(  # noqa: C901, WPS231
                     tasks_set.remove(task_id)
                 except LookupError:
                     continue
-        await asyncio.sleep(check_interval)
+        if tasks_set:
+            await asyncio.sleep(check_interval)
 
     results = []
     for task_id in ordered_ids:  # noqa: WPS440