From b7b152d84d1a7fa1fe4b48f2d98579f4b9c826c3 Mon Sep 17 00:00:00 2001
From: shoucandanghehe <wallfjjd@gmail.com>
Date: Tue, 4 Jun 2024 20:14:19 +0800
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20=E5=AE=8C=E5=96=84=20retry=20?=
 =?UTF-8?q?=E8=A3=85=E9=A5=B0=E5=99=A8=E7=9A=84=E7=B1=BB=E5=9E=8B=EF=BC=8C?=
 =?UTF-8?q?=E5=B9=B6=E6=B7=BB=E5=8A=A0=E6=B6=88=E6=81=AF=E6=8F=90=E7=A4=BA?=
 =?UTF-8?q?=E5=8A=9F=E8=83=BD?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 nonebot_plugin_tetris_stats/utils/retry.py | 30 ++++++++++++----------
 1 file changed, 16 insertions(+), 14 deletions(-)

diff --git a/nonebot_plugin_tetris_stats/utils/retry.py b/nonebot_plugin_tetris_stats/utils/retry.py
index f6b3c98..5ba2874 100644
--- a/nonebot_plugin_tetris_stats/utils/retry.py
+++ b/nonebot_plugin_tetris_stats/utils/retry.py
@@ -1,38 +1,40 @@
 from asyncio import sleep
 from collections.abc import Awaitable, Callable
+from contextlib import suppress
 from datetime import timedelta
 from functools import wraps
-from typing import TypeVar, cast
+from typing import ParamSpec, TypeVar, cast
 
 from nonebot.log import logger
+from nonebot_plugin_alconna.uniseg import SerializeFailed, UniMessage
 
 T = TypeVar('T')
+P = ParamSpec('P')
 
 
 def retry(
     max_attempts: int = 3,
     exception_type: type[BaseException] | tuple[type[BaseException], ...] = Exception,
     delay: timedelta | None = None,
-) -> Callable[[Callable[..., Awaitable[T]]], Callable[..., Awaitable[T]]]:
-    def decorator(func: Callable[..., Awaitable[T]]) -> Callable[..., Awaitable[T]]:
+    reply: str | UniMessage | None = None,
+) -> Callable[[Callable[P, Awaitable[T]]], Callable[P, Awaitable[T]]]:
+    def decorator(func: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T]]:
         @wraps(func)
-        async def wrapper(*args, **kwargs) -> T:  # noqa: ANN002, ANN003
-            attempts = 0
-            while attempts < max_attempts + 1:
+        async def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
+            for i in range(1, max_attempts + 1):
                 try:
                     return await func(*args, **kwargs)
                 except exception_type as e:  # noqa: PERF203
                     logger.exception(e)
-                    attempts += 1
-                    if attempts <= max_attempts:
-                        if delay is not None:
-                            await sleep(delay.total_seconds())
-                        logger.debug(f'Retrying: {func.__name__} ({attempts}/{max_attempts})')
-                        continue
-                    raise
+                    if delay is not None:
+                        await sleep(delay.total_seconds())
+                    message = f'Retrying: {func.__name__} ({i}/{max_attempts})'
+                    logger.debug(message)
+                    with suppress(SerializeFailed):
+                        await UniMessage(reply or message).send()
             msg = 'Unexpectedly reached the end of the retry loop'
             raise RuntimeError(msg)
 
-        return cast(Callable[..., Awaitable[T]], wrapper)
+        return cast(Callable[P, Awaitable[T]], wrapper)
 
     return decorator