[PATCH v2 07/17] telegram: fix: Fix the type annoations for the decorator

[Muhammad Rizki <kiizuha@gnuweeb.org>]

In the previous commit, I created a TODO to fix the typing. In this
commit, I have fixed the decorator's typing.

Fixes: ba02aa3376a4 ("daemon: Add typing in decorator")
Signed-off-by: Muhammad Rizki <kiizuha@gnuweeb.org>
---
 daemon/telegram/packages/decorator.py | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/daemon/telegram/packages/decorator.py b/daemon/telegram/packages/decorator.py
index efcfc38..6bfd8bc 100644
--- a/daemon/telegram/packages/decorator.py
+++ b/daemon/telegram/packages/decorator.py
@@ -4,25 +4,24 @@
 #
 
 from pyrogram.errors.exceptions.flood_420 import FloodWait
-from pyrogram.types import Message
-from typing import Any, Callable, TypeVar
+from typing import Any, Callable, TypeVar, Coroutine
+from typing_extensions import ParamSpec, ParamSpecArgs, ParamSpecKwargs
 from functools import wraps
 import re
 import asyncio
 
 __all__ = ["handle_flood"]
 
-T = TypeVar("T", bound=Message)
+T = TypeVar("T")
+P = ParamSpec("P")
 
-#
-# TODO(Muhammad Rizki): Add more typing for @handle_flood
-#
-def handle_flood(func: Callable[[T], T]) -> Callable[[T], T]:
+
+def handle_flood(func: Callable[P, Coroutine[Any,Any,T]]) -> Callable[P, Coroutine[Any,Any,T]]:
 	@wraps(func)
-	async def callback(*args: Any) -> Any:
+	async def callback(*args: ParamSpecArgs, **kwargs: ParamSpecKwargs) -> T:
 		while True:
 			try:
-				return await func(*args)
+				return await func(*args, **kwargs)
 			except FloodWait as e:
 				# Calling logger attr from the DaemonTelegram() class
 				logger = args[0].logger
-- 
Muhammad Rizki