Merge pull request #18 from taskiq-python/feature/plugins

Plugin system.

Author: s3rius

.flake8

@@ -8,6 +8,8 @@ docstring_style=sphinx
 ignore =
   ; Found `%` string formatting
   WPS323,
+  ; Found block variable overlap.
+  WPS440,
   ; Found multi-line function type annotation
   WPS320,
   ; Found statement that has no effect

taskiq/abc/broker.py

@@ -1,13 +1,34 @@
 from abc import ABC, abstractmethod
 from functools import wraps
 from logging import getLogger
-from typing import Any, AsyncGenerator, Callable, Dict, Optional, Union, overload
+from typing import (  # noqa: WPS235
+    TYPE_CHECKING,
+    Any,
+    AsyncGenerator,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    TypeVar,
+    Union,
+    overload,
+)
+
+from typing_extensions import ParamSpec
 
-from taskiq.abc.result_backend import AsyncResultBackend
 from taskiq.decor import AsyncTaskiqDecoratedTask
-from taskiq.message import TaskiqMessage
+from taskiq.message import BrokerMessage
+from taskiq.plugins.json_formatter import JSONFormatter
 from taskiq.result_backends.dummy import DummyResultBackend
-from taskiq.types_helpers import T_, FuncParams_, ReturnType_
+
+if TYPE_CHECKING:
+    from taskiq.abc.formatter import TaskiqFormatter
+    from taskiq.abc.middleware import TaskiqMiddleware
+    from taskiq.abc.result_backend import AsyncResultBackend
+
+_T = TypeVar("_T")  # noqa: WPS111
+_FuncParams = ParamSpec("_FuncParams")
+_ReturnType = TypeVar("_ReturnType")
 
 logger = getLogger("taskiq")
 
@@ -25,13 +46,28 @@ class AsyncBroker(ABC):
 
     def __init__(
         self,
-        result_backend: Optional[AsyncResultBackend[T_]] = None,
+        result_backend: "Optional[AsyncResultBackend[_T]]" = None,
     ) -> None:
         if result_backend is None:
             result_backend = DummyResultBackend()
+        self.middlewares: "List[TaskiqMiddleware]" = []
         self.result_backend = result_backend
         self.is_worker_process = False
         self.decorator_class = AsyncTaskiqDecoratedTask
+        self.formatter: "TaskiqFormatter" = JSONFormatter()
+
+    def add_middlewares(self, middlewares: "List[TaskiqMiddleware]") -> None:
+        """
+        Add a list of middlewares.
+
+        You should call this method to set middlewares,
+        since it saves current broker in all middlewares.
+
+        :param middlewares: list of middlewares.
+        """
+        for middleware in middlewares:
+            middleware.set_broker(self)
+            self.middlewares.append(middleware)
 
     async def startup(self) -> None:
         """Do something when starting broker."""
@@ -47,7 +83,7 @@ async def shutdown(self) -> None:
     @abstractmethod
     async def kick(
         self,
-        message: TaskiqMessage,
+        message: BrokerMessage,
     ) -> None:
         """
         This method is used to kick tasks out from current program.
@@ -59,7 +95,7 @@ async def kick(
         """
 
     @abstractmethod
-    def listen(self) -> AsyncGenerator[TaskiqMessage, None]:
+    def listen(self) -> AsyncGenerator[BrokerMessage, None]:
         """
         This function listens to new messages and yields them.
 
@@ -73,8 +109,8 @@ def listen(self) -> AsyncGenerator[TaskiqMessage, None]:
     @overload
     def task(
         self,
-        task_name: Callable[FuncParams_, ReturnType_],
-    ) -> AsyncTaskiqDecoratedTask[FuncParams_, ReturnType_]:
+        task_name: Callable[_FuncParams, _ReturnType],
+    ) -> AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType]:
         ...
 
     @overload
@@ -83,8 +119,8 @@ def task(
         task_name: Optional[str] = None,
         **labels: Union[str, int],
     ) -> Callable[
-        [Callable[FuncParams_, ReturnType_]],
-        AsyncTaskiqDecoratedTask[FuncParams_, ReturnType_],
+        [Callable[_FuncParams, _ReturnType]],
+        AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType],
     ]:
         ...
 
@@ -120,12 +156,12 @@ def make_decorated_task(
             inner_labels: Dict[str, Union[str, int]],
             inner_task_name: Optional[str] = None,
         ) -> Callable[
-            [Callable[FuncParams_, ReturnType_]],
-            AsyncTaskiqDecoratedTask[FuncParams_, ReturnType_],
+            [Callable[_FuncParams, _ReturnType]],
+            AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType],
         ]:
             def inner(
-                func: Callable[FuncParams_, ReturnType_],
-            ) -> AsyncTaskiqDecoratedTask[FuncParams_, ReturnType_]:
+                func: Callable[_FuncParams, _ReturnType],
+            ) -> AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType]:
                 nonlocal inner_task_name  # noqa: WPS420
                 if inner_task_name is None:
                     inner_task_name = (  # noqa: WPS442

taskiq/abc/formatter.py

@@ -0,0 +1,27 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict
+
+from taskiq.message import BrokerMessage, TaskiqMessage
+
+
+class TaskiqFormatter(ABC):
+    """Custom formatter for brokers."""
+
+    @abstractmethod
+    def dumps(self, message: TaskiqMessage, labels: Dict[str, Any]) -> BrokerMessage:
+        """
+        Dump message to broker message instance.
+
+        :param message: message to send.
+        :param labels: task's labels.
+        :return: message for brokers.
+        """
+
+    @abstractmethod
+    def loads(self, message: BrokerMessage) -> TaskiqMessage:
+        """
+        Parses broker message to TaskiqMessage.
+
+        :param message: message to parse.
+        :return: parsed taskiq message.
+        """

taskiq/abc/middleware.py

@@ -0,0 +1,85 @@
+from typing import TYPE_CHECKING, Any, Coroutine, Dict, Union
+
+if TYPE_CHECKING:
+    from taskiq.abc.broker import AsyncBroker
+    from taskiq.message import TaskiqMessage
+    from taskiq.result import TaskiqResult
+
+
+class TaskiqMiddleware:
+    """Base class for middlewares."""
+
+    def __init__(self) -> None:
+        self.broker: "AsyncBroker" = None  # type: ignore
+
+    def set_broker(self, broker: "AsyncBroker") -> None:
+        """
+        Sets broker to middleware.
+
+        :param broker: broker to set.
+        """
+        self.broker = broker
+
+    def pre_send(
+        self,
+        message: "TaskiqMessage",
+        labels: Dict[str, Any],
+    ) -> "Union[TaskiqMessage, Coroutine[Any, Any, TaskiqMessage]]":
+        """
+        Hook that executes before sending the task to worker.
+
+        This is a client-side hook, that executes right before
+        the message is sent to broker.
+
+        :param message: message to send.
+        :param labels: task's labels.
+        :return: modified message.
+        """
+        return message
+
+    def post_send(
+        self,
+        message: "TaskiqMessage",
+        labels: Dict[str, Any],
+    ) -> "Union[None, Coroutine[Any, Any, None]]":
+        """
+        This hook is executed right after the task is sent.
+
+        This is a client-side hook. It executes right
+        after the messages is kicked in broker.
+
+        :param message: kicked message.
+        :param labels: labels for a message.
+        """
+
+    def pre_execute(
+        self,
+        message: "TaskiqMessage",
+        labels: Dict[str, Any],
+    ) -> "Union[TaskiqMessage, Coroutine[Any, Any, TaskiqMessage]]":
+        """
+        This hook is called before executing task.
+
+        This is a worker-side hook, wich means it
+        executes in the worker process.
+
+        :param message: incoming parsed taskiq message.
+        :param labels: task's labels without user-supplied lables.
+        :return: modified message.
+        """
+        return message
+
+    def post_execute(
+        self,
+        result: "TaskiqResult[Any]",
+        labels: Dict[str, Any],
+    ) -> "Union[None, Coroutine[Any, Any, None]]":
+        """
+        This hook executes after task is complete.
+
+        This is a worker-side hook. It's called
+        in worker process.
+
+        :param result: result of execution for current task.
+        :param labels: task's labels. Without user-supplied labels.
+        """

taskiq/abc/result_backend.py

@@ -1,12 +1,12 @@
 from abc import ABC, abstractmethod
-from typing import Generic
+from typing import Generic, TypeVar
 
 from taskiq.result import TaskiqResult
-from taskiq.task import AsyncTaskiqTask
-from taskiq.types_helpers import ReturnType_
 
+_ReturnType = TypeVar("_ReturnType")
 
-class AsyncResultBackend(ABC, Generic[ReturnType_]):
+
+class AsyncResultBackend(ABC, Generic[_ReturnType]):
     """Async result backend."""
 
     async def startup(self) -> None:
@@ -15,21 +15,8 @@ async def startup(self) -> None:
     async def shutdown(self) -> None:
         """Do something on shutdown."""
 
-    def generate_task(self, task_id: str) -> "AsyncTaskiqTask[ReturnType_]":
-        """
-        Generates new task.
-
-        This function creates new AsyncTaskiqTask
-        that returned to client after calling kiq
-        method.
-
-        :param task_id: id of a task to save.
-        :return: task object.
-        """
-        return AsyncTaskiqTask(task_id=task_id, result_backend=self)
-
     @abstractmethod
-    async def set_result(self, task_id: str, result: TaskiqResult[ReturnType_]) -> None:
+    async def set_result(self, task_id: str, result: TaskiqResult[_ReturnType]) -> None:
         """
         Saves result to the result backend.
 
@@ -55,7 +42,7 @@ async def get_result(
         self,
         task_id: str,
         with_logs: bool = False,
-    ) -> TaskiqResult[ReturnType_]:
+    ) -> TaskiqResult[_ReturnType]:
         """
         Gets result from the task.
 

taskiq/brokers/inmemory_broker.py

@@ -7,7 +7,7 @@
 from taskiq.abc.result_backend import AsyncResultBackend, TaskiqResult
 from taskiq.cli.async_task_runner import run_task
 from taskiq.exceptions import TaskiqError
-from taskiq.message import TaskiqMessage
+from taskiq.message import BrokerMessage
 
 _ReturnType = TypeVar("_ReturnType")
 
@@ -101,7 +101,7 @@ def __init__(
             logs_format = "%(levelname)s %(message)s"
         self.logs_format = logs_format
 
-    async def kick(self, message: TaskiqMessage) -> None:
+    async def kick(self, message: BrokerMessage) -> None:
         """
         Kicking task.
 
@@ -111,18 +111,19 @@ async def kick(self, message: TaskiqMessage) -> None:
         :raises TaskiqError: if someone wants to kick unknown task.
         """
         target_task = self.available_tasks.get(message.task_name)
+        taskiq_message = self.formatter.loads(message=message)
         if target_task is None:
             raise TaskiqError("Unknown task.")
         result = await run_task(
             target=target_task.original_func,
             signature=inspect.signature(target_task.original_func),
-            message=message,
+            message=taskiq_message,
             log_collector_format=self.logs_format,
             executor=self.executor,
         )
         await self.result_backend.set_result(message.task_id, result)
 
-    async def listen(self) -> AsyncGenerator[TaskiqMessage, None]:  # type: ignore
+    async def listen(self) -> AsyncGenerator[BrokerMessage, None]:  # type: ignore
         """
         Inmemory broker cannot listen.