diff --git a/nonebot_plugin_tetris_stats/utils/retry.py b/nonebot_plugin_tetris_stats/utils/retry.py index 5ba2874..66f5aad 100644 --- a/nonebot_plugin_tetris_stats/utils/retry.py +++ b/nonebot_plugin_tetris_stats/utils/retry.py @@ -1,9 +1,9 @@ from asyncio import sleep -from collections.abc import Awaitable, Callable +from collections.abc import Callable, Coroutine from contextlib import suppress from datetime import timedelta from functools import wraps -from typing import ParamSpec, TypeVar, cast +from typing import Any, ParamSpec, TypeVar from nonebot.log import logger from nonebot_plugin_alconna.uniseg import SerializeFailed, UniMessage @@ -17,24 +17,28 @@ def retry( exception_type: type[BaseException] | tuple[type[BaseException], ...] = Exception, delay: timedelta | None = None, 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]]: +) -> Callable[[Callable[P, Coroutine[Any, Any, T]]], Callable[P, Coroutine[Any, Any, T]]]: + def decorator(func: Callable[P, Coroutine[Any, Any, T]]) -> Callable[P, Coroutine[Any, Any, T]]: @wraps(func) 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) - if delay is not None: - await sleep(delay.total_seconds()) + for i in range(max_attempts + 1): + if i > 0: 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) + if i == max_attempts: + break + try: + return await func(*args, **kwargs) + except exception_type as e: + if i == max_attempts: + raise + logger.exception(e) + if delay is not None: + await sleep(delay.total_seconds()) + return await func(*args, **kwargs) - return cast(Callable[P, Awaitable[T]], wrapper) + return wrapper return decorator