Added kicker object.

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

Author: s3rius

taskiq/abc/broker.py

@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+from dataclasses import asdict, is_dataclass
 from functools import wraps
 from logging import getLogger
 from typing import (  # noqa: WPS235
@@ -29,42 +30,65 @@
 _ReturnType = TypeVar("_ReturnType")
 
 
-class AsyncTaskiqDecoratedTask(Generic[_FuncParams, _ReturnType]):
-    """
-    Class for all task functions.
-
-    When function is decorated
-    with the `task` decorator, it
-    will return an instance of this class.
-
-    This class parametrized with original function's
-    arguments types and a return type.
-
-    This class has kiq method which is used
-    to kick tasks out of this thread and send them to
-    current broker.
-    """
+class AsyncKicker(Generic[_FuncParams, _ReturnType]):
+    """Class that used to modify data before sending it to broker."""
 
     def __init__(
         self,
-        broker: "AsyncBroker",
         task_name: str,
-        original_func: Callable[_FuncParams, _ReturnType],
+        broker: "AsyncBroker",
         labels: Dict[str, Any],
     ) -> None:
-        self.broker = broker
         self.task_name = task_name
-        self.original_func = original_func
+        self.broker = broker
         self.labels = labels
 
-    # Docs for this method are ommited in order to help
-    # your IDE resolve correct docs for it.
-    def __call__(  # noqa: D102
+    def with_label(
         self,
-        *args: _FuncParams.args,
-        **kwargs: _FuncParams.kwargs,
-    ) -> _ReturnType:
-        return self.original_func(*args, **kwargs)
+        label_name: str,
+        value: Any,
+    ) -> "AsyncKicker[_FuncParams, _ReturnType]":
+        """
+        Update one single label.
+
+        This method is used to update
+        task's labels before sending.
+
+        :param label_name: name of the label to update.
+        :param value: label's value.
+        :return: kicker object with new labels.
+        """
+        self.labels[label_name] = value
+        return self
+
+    def with_labels(
+        self,
+        labels: Dict[str, Any],
+    ) -> "AsyncKicker[_FuncParams, _ReturnType]":
+        """
+        Update function's labels before sending.
+
+        :param labels: dict with new labels.
+        :return: kicker with new labels.
+        """
+        self.labels.update(labels)
+        return self
+
+    def with_broker(
+        self,
+        broker: "AsyncBroker",
+    ) -> "AsyncKicker[_FuncParams, _ReturnType]":
+        """
+        Replace broker for the function.
+
+        This method can be used with
+        shared tasks.
+
+        :param broker: new broker instance.
+        :return: Kicker with new broker.
+        """
+        self.broker = broker
+        return self
 
     async def kiq(
         self,
@@ -89,13 +113,21 @@ async def kiq(
         await self.broker.kick(message)
         return self.broker.result_backend.generate_task(message.task_id)  # type: ignore
 
-    def __repr__(self) -> str:
-        return f"AsyncTaskiqDecoratedTask({self.task_name})"
-
     @classmethod
     def _prepare_arg(cls, arg: Any) -> Any:
+        """
+        Parses argument if possible.
+
+        This function is used to construct dicts
+        from pydantic models or dataclasses.
+
+        :param arg: argument to format.
+        :return: Formatted argument.
+        """
         if isinstance(arg, BaseModel):
             arg = arg.dict()
+        if is_dataclass(arg):
+            arg = asdict(arg)
         return arg
 
     def _prepare_message(  # noqa: WPS210
@@ -128,6 +160,80 @@ def _prepare_message(  # noqa: WPS210
         )
 
 
+class AsyncTaskiqDecoratedTask(Generic[_FuncParams, _ReturnType]):
+    """
+    Class for all task functions.
+
+    When function is decorated
+    with the `task` decorator, it
+    will return an instance of this class.
+
+    This class parametrized with original function's
+    arguments types and a return type.
+
+    This class has kiq method which is used
+    to kick tasks out of this thread and send them to
+    current broker.
+    """
+
+    def __init__(
+        self,
+        broker: "AsyncBroker",
+        task_name: str,
+        original_func: Callable[_FuncParams, _ReturnType],
+        labels: Dict[str, Any],
+    ) -> None:
+        self.broker = broker
+        self.task_name = task_name
+        self.original_func = original_func
+        self.labels = labels
+
+    # Docs for this method are ommited in order to help
+    # your IDE resolve correct docs for it.
+    def __call__(  # noqa: D102
+        self,
+        *args: _FuncParams.args,
+        **kwargs: _FuncParams.kwargs,
+    ) -> _ReturnType:
+        return self.original_func(*args, **kwargs)
+
+    async def kiq(
+        self,
+        *args: _FuncParams.args,
+        **kwargs: _FuncParams.kwargs,
+    ) -> AsyncTaskiqTask[_ReturnType]:
+        """
+        This method sends function call over the network.
+
+        It gets current broker and calls it's kick method,
+        returning what it returns.
+
+        :param args: function's arguments.
+        :param kwargs: function's key word arguments.
+
+        :returns: taskiq task.
+        """
+        return await self.kicker().kiq(*args, **kwargs)
+
+    def kicker(self) -> AsyncKicker[_FuncParams, _ReturnType]:
+        """
+        This function returns kicker object.
+
+        Kicker is a object that can modyfy kiq request
+        before sendig it.
+
+        :return: AsyncKicker instance.
+        """
+        return AsyncKicker(
+            task_name=self.task_name,
+            broker=self.broker,
+            labels=self.labels,
+        )
+
+    def __repr__(self) -> str:
+        return f"AsyncTaskiqDecoratedTask({self.task_name})"
+
+
 class AsyncBroker(ABC):
     """
     Async broker.