Added sync methods.

s3rius · s3rius · commit fff31a55b621 · 2022-08-08T22:31:33.000+04:00
Signed-off-by: Pavel Kirilin &lt;win10@list.ru&gt;
diff --git a/taskiq/cli/async_task_runner.py b/taskiq/cli/async_task_runner.py
@@ -16,6 +16,7 @@
 from taskiq.cli.log_collector import log_collector
 from taskiq.message import TaskiqMessage
 from taskiq.result import TaskiqResult
+from taskiq.utils import maybe_awaitable
 
 logger = getLogger("taskiq.worker")
 
@@ -168,13 +169,13 @@ async def run_task(  # noqa: C901, WPS210, WPS211
     )
     if found_exception is not None:
         for middleware in middlewares:
-            err_handler = middleware.on_error(
-                message,
-                result,
-                found_exception,
+            await maybe_awaitable(
+                middleware.on_error(
+                    message,
+                    result,
+                    found_exception,
+                ),
             )
-            if inspect.isawaitable(err_handler):
-                await err_handler
 
     return result
 
@@ -309,13 +310,12 @@ async def async_listen_messages(  # noqa: C901, WPS210, WPS213
             )
             continue
         for middleware in broker.middlewares:
-            pre_ex_res = middleware.pre_execute(
-                taskiq_msg,
+            taskiq_msg = await maybe_awaitable(
+                middleware.pre_execute(
+                    taskiq_msg,
+                ),
             )
-            if inspect.isawaitable(pre_ex_res):
-                taskiq_msg = await pre_ex_res
-            else:
-                taskiq_msg = pre_ex_res  # type: ignore
+
         result = await run_task(
             target=broker.available_tasks[message.task_name].original_func,
             signature=task_signatures.get(message.task_name),
@@ -325,9 +325,7 @@ async def async_listen_messages(  # noqa: C901, WPS210, WPS213
             middlewares=broker.middlewares,
         )
         for middleware in broker.middlewares:
-            post_ex_res = middleware.post_execute(taskiq_msg, result)
-            if inspect.isawaitable(post_ex_res):
-                await post_ex_res
+            await maybe_awaitable(middleware.post_execute(taskiq_msg, result))
         try:
             await broker.result_backend.set_result(message.task_id, result)
         except Exception as exc:
diff --git a/taskiq/decor.py b/taskiq/decor.py
@@ -12,7 +12,7 @@
 from typing_extensions import ParamSpec
 
 from taskiq.kicker import AsyncKicker
-from taskiq.task import AsyncTaskiqTask
+from taskiq.task import AsyncTaskiqTask, SyncTaskiqTask
 
 if TYPE_CHECKING:
     from taskiq.abc.broker import AsyncBroker
@@ -93,6 +93,40 @@ async def kiq(
         """
         return await self.kicker().kiq(*args, **kwargs)
 
+    @overload
+    def kiq_sync(
+        self: "AsyncTaskiqDecoratedTask[_FuncParams, Coroutine[Any, Any, _T]]",
+        *args: _FuncParams.args,
+        **kwargs: _FuncParams.kwargs,
+    ) -> SyncTaskiqTask[_T]:
+        ...
+
+    @overload
+    def kiq_sync(
+        self: "AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType]",
+        *args: _FuncParams.args,
+        **kwargs: _FuncParams.kwargs,
+    ) -> SyncTaskiqTask[_ReturnType]:
+        ...
+
+    def kiq_sync(
+        self,
+        *args: _FuncParams.args,
+        **kwargs: _FuncParams.kwargs,
+    ) -> Any:
+        """
+        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 self.kicker().kiq_sync(*args, **kwargs)
+
     def kicker(self) -> AsyncKicker[_FuncParams, _ReturnType]:
         """
         This function returns kicker object.
diff --git a/taskiq/kicker.py b/taskiq/kicker.py
@@ -1,5 +1,4 @@
 from dataclasses import asdict, is_dataclass
-from inspect import isawaitable
 from logging import getLogger
 from typing import (
     TYPE_CHECKING,
@@ -17,7 +16,8 @@
 
 from taskiq.exceptions import SendTaskError
 from taskiq.message import TaskiqMessage
-from taskiq.task import AsyncTaskiqTask
+from taskiq.task import AsyncTaskiqTask, SyncTaskiqTask
+from taskiq.utils import maybe_awaitable, run_sync
 
 if TYPE_CHECKING:
     from taskiq.abc.broker import AsyncBroker
@@ -87,7 +87,7 @@ async def kiq(  # noqa: D102
     ) -> AsyncTaskiqTask[_ReturnType]:
         ...
 
-    async def kiq(  # noqa: C901
+    async def kiq(
         self,
         *args: _FuncParams.args,
         **kwargs: _FuncParams.kwargs,
@@ -110,24 +110,55 @@ async def kiq(  # noqa: C901
         )
         message = self._prepare_message(*args, **kwargs)
         for middleware in self.broker.middlewares:
-            pre_send_res = middleware.pre_send(message)
-            if isawaitable(pre_send_res):
-                message = await pre_send_res
-            else:
-                message = pre_send_res  # type: ignore
+            message = await maybe_awaitable(middleware.pre_send(message))
+
         try:
             await self.broker.kick(self.broker.formatter.dumps(message))
         except Exception as exc:
             raise SendTaskError() from exc
+
         for middleware in self.broker.middlewares:
-            post_send_res = middleware.post_send(message)
-            if isawaitable(post_send_res):
-                await post_send_res
+            await maybe_awaitable(middleware.post_send(message))
+
         return AsyncTaskiqTask(
             task_id=message.task_id,
             result_backend=self.broker.result_backend,
         )
 
+    @overload
+    def kiq_sync(
+        self: "AsyncKicker[_FuncParams, Coroutine[Any, Any, _T]]",
+        *args: _FuncParams.args,
+        **kwargs: _FuncParams.kwargs,
+    ) -> SyncTaskiqTask[_T]:
+        ...
+
+    @overload
+    def kiq_sync(
+        self: "AsyncKicker[_FuncParams, _ReturnType]",
+        *args: _FuncParams.args,
+        **kwargs: _FuncParams.kwargs,
+    ) -> SyncTaskiqTask[_ReturnType]:
+        ...
+
+    def kiq_sync(
+        self,
+        *args: _FuncParams.args,
+        **kwargs: _FuncParams.kwargs,
+    ) -> Any:
+        """
+        This method sends function call over the network.
+
+        It just wraps async kiq call in run_sync
+        funcion.
+
+        :param args: function's arguments.
+        :param kwargs: function's key word arguments.
+
+        :returns: sync taskiq task.
+        """
+        return SyncTaskiqTask(run_sync(self.kiq(*args, **kwargs)))
+
     @classmethod
     def _prepare_arg(cls, arg: Any) -> Any:
         """
diff --git a/taskiq/task.py b/taskiq/task.py
@@ -1,12 +1,14 @@
 import asyncio
+from abc import ABC, abstractmethod
 from time import time
-from typing import TYPE_CHECKING, Generic, TypeVar
+from typing import TYPE_CHECKING, Any, Coroutine, Generic, TypeVar, Union
 
 from taskiq.exceptions import (
     ResultGetError,
     ResultIsReadyError,
     TaskiqResultTimeoutError,
 )
+from taskiq.utils import run_sync
 
 if TYPE_CHECKING:
     from taskiq.abc.result_backend import AsyncResultBackend
@@ -15,7 +17,111 @@
 _ReturnType = TypeVar("_ReturnType")
 
 
-class AsyncTaskiqTask(Generic[_ReturnType]):
+class _Task(ABC, Generic[_ReturnType]):
+    """TaskiqTask interface."""
+
+    @abstractmethod
+    def is_ready(self) -> Union[bool, Coroutine[Any, Any, bool]]:
+        """
+        Method to check wether result is ready.
+
+        :return: True if result is ready.
+        """
+
+    @abstractmethod
+    def get_result(  # noqa: WPS234
+        self,
+        with_logs: bool = False,
+    ) -> Union[
+        "TaskiqResult[_ReturnType]",
+        Coroutine[Any, Any, "TaskiqResult[_ReturnType]"],
+    ]:
+        """
+        Get actual execution result.
+
+        :param with_logs: wether you want to fetch logs.
+        :return: TaskiqResult.
+        """
+
+    @abstractmethod
+    def wait_result(  # noqa: WPS234
+        self,
+        check_interval: float = 0.2,
+        timeout: float = -1.0,
+        with_logs: bool = False,
+    ) -> Union[
+        "TaskiqResult[_ReturnType]",
+        Coroutine[Any, Any, "TaskiqResult[_ReturnType]"],
+    ]:
+        """
+        Wait for result to become ready and get it.
+
+        This function constantly checks wheter result is ready
+        and fetches it when it becomes available.
+
+        :param check_interval: how ofen availability is checked.
+        :param timeout: maximum amount of time it will wait
+            before raising TaskiqResultTimeoutError.
+        :param with_logs: whether you need to download logs.
+        :return: TaskiqResult.
+        """
+
+
+class SyncTaskiqTask(_Task[_ReturnType]):
+    """Sync wrapper over AsyncTaskiqTask."""
+
+    def __init__(self, async_task: "AsyncTaskiqTask[_ReturnType]") -> None:
+        self.async_task = async_task
+
+    def is_ready(self) -> bool:
+        """
+        Checks if task is completed.
+
+        :return: True if task is completed.
+        """
+        return run_sync(self.async_task.is_ready())
+
+    def get_result(self, with_logs: bool = False) -> "TaskiqResult[_ReturnType]":
+        """
+        Get result of a task from result backend.
+
+        :param with_logs: whether you want to fetch logs from worker.
+
+        :return: task's return value.
+        """
+        return run_sync(self.async_task.get_result(with_logs=with_logs))
+
+    def wait_result(
+        self,
+        check_interval: float = 0.2,
+        timeout: float = -1,
+        with_logs: bool = False,
+    ) -> "TaskiqResult[_ReturnType]":
+        """
+        Waits until result is ready.
+
+        This method just checks whether the task is
+        ready. And if it is it returns the result.
+
+        It may throw TaskiqResultTimeoutError if
+        task didn't became ready in provided
+        period of time.
+
+        :param check_interval: How often checks are performed.
+        :param timeout: timeout for the result.
+        :param with_logs: whether you want to fetch logs from worker.
+        :return: task's return value.
+        """
+        return run_sync(
+            self.async_task.wait_result(
+                check_interval=check_interval,
+                timeout=timeout,
+                with_logs=with_logs,
+            ),
+        )
+
+
+class AsyncTaskiqTask(_Task[_ReturnType]):
     """AsyncTask for AsyncResultBackend."""
 
     def __init__(
@@ -59,8 +165,8 @@ async def get_result(self, with_logs: bool = False) -> "TaskiqResult[_ReturnType
 
     async def wait_result(
         self,
-        check_interval: float = 1.0,
-        timeout: float = 5.0,
+        check_interval: float = 0.2,
+        timeout: float = -1.0,
         with_logs: bool = False,
     ) -> "TaskiqResult[_ReturnType]":
         """
@@ -83,6 +189,6 @@ async def wait_result(
         start_time = time()
         while not await self.is_ready():
             await asyncio.sleep(check_interval)
-            if time() - start_time > timeout:
+            if 0 < timeout < time() - start_time:
                 raise TaskiqResultTimeoutError()
         return await self.get_result(with_logs=with_logs)
diff --git a/taskiq/utils.py b/taskiq/utils.py
@@ -0,0 +1,53 @@
+import asyncio
+import inspect
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, Awaitable, Coroutine, TypeVar, Union
+
+_T = TypeVar("_T")  # noqa: WPS111
+
+
+def run_sync(coroutine: "Coroutine[Any, Any, _T]") -> _T:
+    """
+    Run the coroutine synchronously.
+
+    This function tries to run corouting using asyncio.run.
+
+    If it's not possible, it manually creates executor and
+    runs async function returns it's result.
+
+     1. When called within a coroutine.
+     2. When called from ``python -m asyncio``, or iPython with %autoawait
+        enabled, which means an event loop may already be running in the
+        current thread.
+
+    :param coroutine: awaitable to execute.
+    :returns: the same type as if it were awaited.
+    """
+    try:
+        # We try this first, as in most situations this will work.
+        return asyncio.run(coroutine)
+    except RuntimeError:
+        # An event loop already exists.
+        with ThreadPoolExecutor(max_workers=1) as executor:
+            future = executor.submit(asyncio.run, coroutine)
+            return future.result()
+
+
+async def maybe_awaitable(
+    possible_coroutine: "Union[_T, Awaitable[_T]]",
+) -> _T:
+    """
+    Awaits coroutine if needed.
+
+    This function allows run function
+    that may return coroutine.
+
+    It not awaitable value passed, it
+    returned immediately.
+
+    :param possible_coroutine: some value.
+    :return: value.
+    """
+    if inspect.isawaitable(possible_coroutine):
+        return await possible_coroutine
+    return possible_coroutine  # type: ignore
