Added custom formatters. Fixed mypy.

Signed-off-by: Pavel Kirilin <[email protected]>

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,29 @@
 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 (
+    Any,
+    AsyncGenerator,
+    Callable,
+    Dict,
+    Optional,
+    TypeVar,
+    Union,
+    overload,
+)
+
+from typing_extensions import ParamSpec
+
+from taskiq.abc.plugins.formatter import TaskiqFormatter
 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_
+
+_T = TypeVar("_T")  # noqa: WPS111
+_FuncParams = ParamSpec("_FuncParams")
+_ReturnType = TypeVar("_ReturnType")
 
 logger = getLogger("taskiq")
 
@@ -25,13 +41,14 @@ 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.result_backend = result_backend
         self.is_worker_process = False
         self.decorator_class = AsyncTaskiqDecoratedTask
+        self.formatter: TaskiqFormatter = JSONFormatter()
 
     async def startup(self) -> None:
         """Do something when starting broker."""
@@ -47,7 +64,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 +76,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 +90,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 +100,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 +137,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/plugins/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/result_backend.py

@@ -1,12 +1,13 @@
 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,7 +16,7 @@ async def startup(self) -> None:
     async def shutdown(self) -> None:
         """Do something on shutdown."""
 
-    def generate_task(self, task_id: str) -> "AsyncTaskiqTask[ReturnType_]":
+    def generate_task(self, task_id: str) -> "AsyncTaskiqTask[_ReturnType]":
         """
         Generates new task.
 
@@ -29,7 +30,7 @@ def generate_task(self, task_id: str) -> "AsyncTaskiqTask[ReturnType_]":
         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 +56,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.
 

taskiq/brokers/shared_broker.py

@@ -4,16 +4,16 @@
 from taskiq.decor import AsyncTaskiqDecoratedTask
 from taskiq.exceptions import TaskiqError
 from taskiq.kicker import AsyncKicker
-from taskiq.message import TaskiqMessage
-from taskiq.types_helpers import ReturnType_
+from taskiq.message import BrokerMessage
 
-Params_ = TypeVar("Params_")  # noqa: WPS120
+_ReturnType = TypeVar("_ReturnType")
+_Params = TypeVar("_Params")
 
 
-class SharedDecoratedTask(AsyncTaskiqDecoratedTask[Params_, ReturnType_]):
+class SharedDecoratedTask(AsyncTaskiqDecoratedTask[_Params, _ReturnType]):
     """Decorator that is used with shared broker."""
 
-    def kicker(self) -> AsyncKicker[Params_, ReturnType_]:
+    def kicker(self) -> AsyncKicker[_Params, _ReturnType]:
         """
         This method updates getting default kicker.
 
@@ -45,7 +45,7 @@ def __init__(self) -> None:
         self._default_broker: Optional[AsyncBroker] = None
         self.decorator_class = SharedDecoratedTask
 
-    async def kick(self, message: TaskiqMessage) -> None:
+    async def kick(self, message: BrokerMessage) -> None:
         """
         Shared broker cannot kick tasks.
 
@@ -62,7 +62,7 @@ def default_broker(self, new_broker: AsyncBroker) -> None:
         """
         self._default_broker = new_broker
 
-    async def listen(self) -> AsyncGenerator[TaskiqMessage, None]:  # type: ignore
+    async def listen(self) -> AsyncGenerator[BrokerMessage, None]:  # type: ignore
         """
         Shared broker cannot listen to tasks.
 

taskiq/cli/async_task_runner.py

@@ -79,7 +79,7 @@ def parse_params(  # noqa: C901
                 continue
             try:
                 message.args[argnum] = parse_obj_as(annot, value)
-            except (ValueError, RuntimeError) as exc:  # noqa: WPS440
+            except (ValueError, RuntimeError) as exc:
                 logger.debug(exc, exc_info=True)
 
 
@@ -266,19 +266,28 @@ async def async_listen_messages(  # noqa: C901, WPS210, WPS213
             "Function for task %s is resolved. Executing...",
             message.task_name,
         )
+        try:
+            taskiq_msg = broker.formatter.loads(message=message)
+        except Exception as exc:
+            logger.warning(
+                "Cannot parse message: %s. Skipping execution.\n %s",
+                message,
+                exc,
+                exc_info=True,
+            )
+            continue
         result = await run_task(
             broker.available_tasks[message.task_name].original_func,
             task_signatures.get(message.task_name),
-            message,
+            taskiq_msg,
             cli_args.log_collector_format,
             executor,
         )
         try:
             await broker.result_backend.set_result(message.task_id, result)
         except Exception as exc:
             logger.exception(
-                "Can't set result in %s result backend: \n%s",
-                broker.result_backend.__class__.__name__,
+                "Can't set result in result backend. Cause: %s",
                 exc,
                 exc_info=True,
             )