From 820b7c6075d3ba530aa9f73ad3bc6793c2868873 Mon Sep 17 00:00:00 2001
From: DABND19 <grigorenkodavid@gmail.com>
Date: Fri, 3 May 2024 21:55:27 +0300
Subject: [PATCH] feat: Added typehints for asyncbackoff and asyncretry.

---
 aiomisc/backoff.py    | 20 ++++++++++++--------
 aiomisc/timeout.py    | 21 +++++++++++++++------
 tests/test_backoff.py |  2 +-
 3 files changed, 28 insertions(+), 15 deletions(-)

diff --git a/aiomisc/backoff.py b/aiomisc/backoff.py
index 97d9c783..203e3aa5 100644
--- a/aiomisc/backoff.py
+++ b/aiomisc/backoff.py
@@ -1,4 +1,5 @@
 import asyncio
+import sys
 from functools import wraps
 from typing import (
     Any, Awaitable, Callable, Optional, Tuple, Type, TypeVar, Union,
@@ -8,12 +9,15 @@
 from .timeout import timeout
 
 
-Number = Union[int, float]
-T = TypeVar("T")
+if sys.version_info >= (3, 10):
+    from typing import ParamSpec
+else:
+    from typing_extensions import ParamSpec
 
 
-WrapReturnType = Callable[..., Awaitable[T]]
-ReturnType = Callable[..., WrapReturnType]
+Number = Union[int, float]
+T = TypeVar("T")
+P = ParamSpec("P")
 
 
 class BackoffStatistic(Statistic):
@@ -38,7 +42,7 @@ def asyncbackoff(
     giveup: Optional[Callable[[Exception], bool]] = None,
     statistic_name: Optional[str] = None,
     statistic_class: Type[BackoffStatistic] = BackoffStatistic,
-) -> ReturnType:
+) -> Callable[[Callable[P, Awaitable[T]]], Callable[P, Awaitable[T]]]:
     """
     Patametric decorator that ensures that ``attempt_timeout`` and
     ``deadline`` time limits are met by decorated function.
@@ -81,12 +85,12 @@ def asyncbackoff(
     exceptions = tuple(exceptions) or ()
     exceptions += asyncio.TimeoutError,
 
-    def decorator(func: WrapReturnType) -> WrapReturnType:
+    def decorator(func: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T]]:
         if attempt_timeout is not None:
             func = timeout(attempt_timeout)(func)
 
         @wraps(func)
-        async def wrap(*args: Any, **kwargs: Any) -> T:
+        async def wrap(*args: P.args, **kwargs: P.kwargs) -> T:
             last_exc = None
             tries = 0
 
@@ -141,7 +145,7 @@ def asyncretry(
     pause: Number = 0,
     giveup: Optional[Callable[[Exception], bool]] = None,
     statistic_name: Optional[str] = None,
-) -> ReturnType:
+) -> Callable[[Callable[P, Awaitable[T]]], Callable[P, Awaitable[T]]]:
     """
     Shortcut of ``asyncbackoff(None, None, 0, Exception)``.
 
diff --git a/aiomisc/timeout.py b/aiomisc/timeout.py
index 330f599a..caa349ad 100644
--- a/aiomisc/timeout.py
+++ b/aiomisc/timeout.py
@@ -1,22 +1,31 @@
 import asyncio
+import sys
 from functools import wraps
-from typing import Any, Awaitable, Callable, TypeVar, Union
+from typing import Awaitable, Callable, TypeVar, Union
+
+
+if sys.version_info >= (3, 10):
+    from typing import ParamSpec
+else:
+    from typing_extensions import ParamSpec
 
 
 T = TypeVar("T")
+P = ParamSpec("P")
 Number = Union[int, float]
-FuncType = Callable[..., Awaitable[T]]
 
 
-def timeout(value: Number) -> Callable[[FuncType], FuncType]:
+def timeout(
+    value: Number
+) -> Callable[[Callable[P, Awaitable[T]]], Callable[P, Awaitable[T]]]:
     def decorator(
-        func: FuncType,
-    ) -> FuncType:
+        func: Callable[P, Awaitable[T]],
+    ) -> Callable[P, Awaitable[T]]:
         if not asyncio.iscoroutinefunction(func):
             raise TypeError("Function is not a coroutine function")
 
         @wraps(func)
-        async def wrap(*args: Any, **kwargs: Any) -> T:
+        async def wrap(*args: P.args, **kwargs: P.kwargs) -> T:
             return await asyncio.wait_for(
                 func(*args, **kwargs),
                 timeout=value,
diff --git a/tests/test_backoff.py b/tests/test_backoff.py
index 0057987e..4de42659 100644
--- a/tests/test_backoff.py
+++ b/tests/test_backoff.py
@@ -219,7 +219,7 @@ def test_values(event_loop):
         aiomisc.asyncbackoff(0, 0, -0.1)
 
     with pytest.raises(TypeError):
-        aiomisc.asyncbackoff(0, 0)(lambda x: None)
+        aiomisc.asyncbackoff(0, 0)(lambda x: None)  # type: ignore
 
 
 async def test_too_long_multiple(event_loop):