nonebot-plugin-tetris-stats/nonebot_plugin_tetris_stats/utils/retry.py

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 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,
    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: 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())
                    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[P, Awaitable[T]], wrapper)

    return decorator