Merge pull request #14 from taskiq-python/feature/pydantic-result

Taskiq result is now a generic model.

Author: s3rius

taskiq/abc/broker.py

@@ -6,6 +6,7 @@
     Any,
     AsyncGenerator,
     Callable,
+    Coroutine,
     Dict,
     Generic,
     Optional,
@@ -90,11 +91,27 @@ def with_broker(
         self.broker = broker
         return self
 
+    @overload
+    async def kiq(  # noqa: D102
+        self: "AsyncKicker[_FuncParams, Coroutine[Any, Any, _T]]",
+        *args: _FuncParams.args,
+        **kwargs: _FuncParams.kwargs,
+    ) -> AsyncTaskiqTask[_T]:
+        ...
+
+    @overload
+    async def kiq(  # noqa: D102
+        self: "AsyncKicker[_FuncParams, _ReturnType]",
+        *args: _FuncParams.args,
+        **kwargs: _FuncParams.kwargs,
+    ) -> AsyncTaskiqTask[_ReturnType]:
+        ...
+
     async def kiq(
         self,
         *args: _FuncParams.args,
         **kwargs: _FuncParams.kwargs,
-    ) -> AsyncTaskiqTask[_ReturnType]:
+    ) -> Any:
         """
         This method sends function call over the network.
 
@@ -111,7 +128,7 @@ async def kiq(
         )
         message = self._prepare_message(*args, **kwargs)
         await self.broker.kick(message)
-        return self.broker.result_backend.generate_task(message.task_id)  # type: ignore
+        return self.broker.result_backend.generate_task(message.task_id)
 
     @classmethod
     def _prepare_arg(cls, arg: Any) -> Any:
@@ -197,11 +214,27 @@ def __call__(  # noqa: D102
     ) -> _ReturnType:
         return self.original_func(*args, **kwargs)
 
+    @overload
+    async def kiq(  # noqa: D102
+        self: "AsyncTaskiqDecoratedTask[_FuncParams, Coroutine[Any, Any, _T]]",
+        *args: _FuncParams.args,
+        **kwargs: _FuncParams.kwargs,
+    ) -> AsyncTaskiqTask[_T]:
+        ...
+
+    @overload
+    async def kiq(  # noqa: D102
+        self: "AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType]",
+        *args: _FuncParams.args,
+        **kwargs: _FuncParams.kwargs,
+    ) -> AsyncTaskiqTask[_ReturnType]:
+        ...
+
     async def kiq(
         self,
         *args: _FuncParams.args,
         **kwargs: _FuncParams.kwargs,
-    ) -> AsyncTaskiqTask[_ReturnType]:
+    ) -> Any:
         """
         This method sends function call over the network.
 

taskiq/abc/result_backend.py

@@ -1,44 +1,23 @@
 import asyncio
 from abc import ABC, abstractmethod
 from time import time
-from typing import Any, Coroutine, Generic, Optional, TypeVar, overload
+from typing import Generic, Optional, TypeVar
+
+from pydantic.generics import GenericModel
 
 from taskiq.exceptions import TaskiqResultTimeoutError
 
 _T = TypeVar("_T")  # noqa: WPS111
 _ReturnType = TypeVar("_ReturnType")
 
 
-class TaskiqResult(Generic[_ReturnType]):
+class TaskiqResult(GenericModel, Generic[_ReturnType]):
     """Result of a remote task invocation."""
 
-    def __init__(
-        self,
-        is_err: bool,
-        log: Optional[str],
-        return_value: Any,
-        execution_time: float,
-    ) -> None:
-        self.is_err = is_err
-        self.log = log
-        self._return_value = return_value
-        self.execution_time = execution_time
-
-    @overload
-    def value(self: "TaskiqResult[Coroutine[Any, Any, _T]]") -> _T:
-        ...
-
-    @overload
-    def value(self) -> _ReturnType:
-        ...
-
-    def value(self) -> Any:
-        """
-        Returns function's return value.
-
-        :return: function's value.
-        """
-        return self._return_value
+    is_err: bool
+    log: Optional[str]
+    return_value: _ReturnType
+    execution_time: float
 
 
 class AsyncResultBackend(ABC, Generic[_ReturnType]):