From e7779b9b2fcdd380db2f7f2ce264c59521da0a3f Mon Sep 17 00:00:00 2001
From: DABND19 <grigorenkodavid@gmail.com>
Date: Tue, 7 May 2024 18:57:29 +0300
Subject: [PATCH] fix: Return coroutines instead of awaitables.

---
 aiomisc/aggregate.py | 15 ++++++++++-----
 aiomisc/backoff.py   | 16 ++++++++++++----
 aiomisc/timeout.py   | 11 +++++++----
 3 files changed, 29 insertions(+), 13 deletions(-)

diff --git a/aiomisc/aggregate.py b/aiomisc/aggregate.py
index 34d15f09..98da210e 100644
--- a/aiomisc/aggregate.py
+++ b/aiomisc/aggregate.py
@@ -7,8 +7,8 @@
 from inspect import Parameter
 from typing import (
     Any,
-    Awaitable,
     Callable,
+    Coroutine,
     Generic,
     Iterable,
     List,
@@ -246,7 +246,7 @@ def __init__(
 
 def aggregate(
     leeway_ms: float, max_count: Optional[int] = None
-) -> Callable[[AggregateFunc[V, R]], Callable[[V], Awaitable[R]]]:
+) -> Callable[[AggregateFunc[V, R]], Callable[[V], Coroutine[Any, Any, R]]]:
     """
     Parametric decorator that aggregates multiple
     (but no more than ``max_count`` defaulting to ``None``) single-argument
@@ -275,7 +275,9 @@ def aggregate(
 
     :return:
     """
-    def decorator(func: AggregateFunc[V, R]) -> Callable[[V], Awaitable[R]]:
+    def decorator(
+        func: AggregateFunc[V, R]
+    ) -> Callable[[V], Coroutine[Any, Any, R]]:
         aggregator = Aggregator(
             func, max_count=max_count, leeway_ms=leeway_ms,
         )
@@ -285,7 +287,10 @@ def decorator(func: AggregateFunc[V, R]) -> Callable[[V], Awaitable[R]]:
 
 def aggregate_async(
     leeway_ms: float, max_count: Optional[int] = None,
-) -> Callable[[AggregateAsyncFunc[V, R]], Callable[[V], Awaitable[R]]]:
+) -> Callable[
+    [AggregateAsyncFunc[V, R]],
+    Callable[[V], Coroutine[Any, Any, R]]
+]:
     """
     Same as ``aggregate``, but with ``func`` arguments of type ``Arg``
     containing ``value`` and ``future`` attributes instead. In this setting
@@ -298,7 +303,7 @@ def aggregate_async(
     """
     def decorator(
         func: AggregateAsyncFunc[V, R]
-    ) -> Callable[[V], Awaitable[R]]:
+    ) -> Callable[[V], Coroutine[Any, Any, R]]:
         aggregator = AggregatorAsync(
             func, max_count=max_count, leeway_ms=leeway_ms,
         )
diff --git a/aiomisc/backoff.py b/aiomisc/backoff.py
index 203e3aa5..7df76720 100644
--- a/aiomisc/backoff.py
+++ b/aiomisc/backoff.py
@@ -2,7 +2,7 @@
 import sys
 from functools import wraps
 from typing import (
-    Any, Awaitable, Callable, Optional, Tuple, Type, TypeVar, Union,
+    Any, Callable, Coroutine, Optional, Tuple, Type, TypeVar, Union,
 )
 
 from .counters import Statistic
@@ -42,7 +42,10 @@ def asyncbackoff(
     giveup: Optional[Callable[[Exception], bool]] = None,
     statistic_name: Optional[str] = None,
     statistic_class: Type[BackoffStatistic] = BackoffStatistic,
-) -> Callable[[Callable[P, Awaitable[T]]], Callable[P, Awaitable[T]]]:
+) -> Callable[
+    [Callable[P, Coroutine[Any, Any, T]]],
+    Callable[P, Coroutine[Any, Any, T]],
+]:
     """
     Patametric decorator that ensures that ``attempt_timeout`` and
     ``deadline`` time limits are met by decorated function.
@@ -85,7 +88,9 @@ def asyncbackoff(
     exceptions = tuple(exceptions) or ()
     exceptions += asyncio.TimeoutError,
 
-    def decorator(func: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T]]:
+    def decorator(
+        func: Callable[P, Coroutine[Any, Any, T]]
+    ) -> Callable[P, Coroutine[Any, Any, T]]:
         if attempt_timeout is not None:
             func = timeout(attempt_timeout)(func)
 
@@ -145,7 +150,10 @@ def asyncretry(
     pause: Number = 0,
     giveup: Optional[Callable[[Exception], bool]] = None,
     statistic_name: Optional[str] = None,
-) -> Callable[[Callable[P, Awaitable[T]]], Callable[P, Awaitable[T]]]:
+) -> Callable[
+    [Callable[P, Coroutine[Any, Any, T]]],
+    Callable[P, Coroutine[Any, Any, T]],
+]:
     """
     Shortcut of ``asyncbackoff(None, None, 0, Exception)``.
 
diff --git a/aiomisc/timeout.py b/aiomisc/timeout.py
index caa349ad..9448f253 100644
--- a/aiomisc/timeout.py
+++ b/aiomisc/timeout.py
@@ -1,7 +1,7 @@
 import asyncio
 import sys
 from functools import wraps
-from typing import Awaitable, Callable, TypeVar, Union
+from typing import Any, Callable, Coroutine, TypeVar, Union
 
 
 if sys.version_info >= (3, 10):
@@ -17,10 +17,13 @@
 
 def timeout(
     value: Number
-) -> Callable[[Callable[P, Awaitable[T]]], Callable[P, Awaitable[T]]]:
+) -> Callable[
+    [Callable[P, Coroutine[Any, Any, T]]],
+    Callable[P, Coroutine[Any, Any, T]],
+]:
     def decorator(
-        func: Callable[P, Awaitable[T]],
-    ) -> Callable[P, Awaitable[T]]:
+        func: Callable[P, Coroutine[Any, Any, T]],
+    ) -> Callable[P, Coroutine[Any, Any, T]]:
         if not asyncio.iscoroutinefunction(func):
             raise TypeError("Function is not a coroutine function")