Source code for pyqqq.utils.retry

import time
from dataclasses import dataclass
from functools import wraps
from typing import Optional

from pyqqq.utils.logger import get_logger




@dataclass
class RetryConfig:
    """retry 데코레이터의 동작을 제어하는 설정 객체입니다.

    설정되지 않은 필드(None)는 데코레이터에 지정된 기본값을 그대로 사용합니다.

    Attributes:
        total_tries: 시도할 총 횟수.
        delay: 재시도 간 초기 지연 시간(초).
        backoff: 백오프 인자; 예를 들어, 2는 각 재시도 간 지연 시간을 두 배로 늘립니다.
        silently: True일 경우 예외 메시지를 출력하지 않습니다.
    """

    total_tries: Optional[int] = None
    delay: Optional[float] = None
    backoff: Optional[float] = None
    silently: Optional[bool] = None



_global_retry_config = RetryConfig()




def configure_retry(
    total_tries: Optional[int] = None,
    delay: Optional[float] = None,
    backoff: Optional[float] = None,
    silently: Optional[bool] = None,
):
    """전역 retry 설정을 변경합니다.

    호출 시 전달된 인자만 덮어쓰며, 전달되지 않은 인자는 기존 값을 유지합니다.
    설정된 값은 이후 모든 @retry 데코레이터 호출에 적용됩니다.

    Args:
        total_tries: 시도할 총 횟수.
        delay: 재시도 간 초기 지연 시간(초).
        backoff: 백오프 인자.
        silently: True일 경우 예외 메시지를 출력하지 않습니다.

    Examples:
        >>> from pyqqq.utils.retry import configure_retry
        >>> configure_retry(total_tries=10, delay=1)
    """
    global _global_retry_config
    if total_tries is not None:
        _global_retry_config.total_tries = total_tries
    if delay is not None:
        _global_retry_config.delay = delay
    if backoff is not None:
        _global_retry_config.backoff = backoff
    if silently is not None:
        _global_retry_config.silently = silently





def reset_retry_config():
    """전역 retry 설정을 초기 상태(모두 None)로 되돌립니다."""
    global _global_retry_config
    _global_retry_config = RetryConfig()





def retry(exceptions, total_tries=5, delay=0.5, backoff=2, silently: bool = False):
    """
    지정된 예외가 발생할 때 함수 실행을 지정된 횟수만큼 재시도하는 데코레이터입니다.

    이 데코레이터는 함수를 실행 중 예외가 발생했을 때, 주어진 조건에 따라 자동으로 함수를 재시도합니다.
    재시도 간격은 지연시간과 백오프 인자에 의해 결정되며, 모든 시도가 실패하면 최종적으로 예외를 발생시킵니다.

    전역 설정(configure_retry)이 있으면 데코레이터 인자보다 우선 적용됩니다.

    Args:
        exceptions (Exception or tuple): 재시도를 트리거하는 예외 또는 예외 튜플.
        total_tries (int): 시도할 총 횟수.
        delay (float): 재시도 간 초기 지연 시간(초).
        backoff (float): 백오프 인자; 예를 들어, 2는 각 재시도 간 지연 시간을 두 배로 늘립니다.
        silently (bool): True일 경우 예외 메시지를 출력하지 않습니다.

    Returns:
        function: 예외 발생 시 재시도를 수행하는 함수.

    Examples:
        >>> @retry(Exception, total_tries=3, delay=1, backoff=1, silently=False)
        ... def test_func():
        ...     print("Trying...")
        ...     raise Exception("An error occurred")
    """
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            logger = get_logger(f"{__name__}.{func.__name__}")

            cfg = _global_retry_config
            _total_tries = cfg.total_tries if cfg.total_tries is not None else total_tries
            _delay = cfg.delay if cfg.delay is not None else delay
            _backoff = cfg.backoff if cfg.backoff is not None else backoff
            _silently = cfg.silently if cfg.silently is not None else silently

            current_try = 0
            while current_try < _total_tries:
                try:
                    return func(*args, **kwargs)
                except exceptions as e:
                    current_try += 1
                    sleep_time = _delay * (_backoff ** (current_try - 1))
                    if current_try != _total_tries:
                        if not _silently:
                            logger.warning(f"{str(e)}\nRetrying in {sleep_time} seconds...")
                        time.sleep(sleep_time)
                    else:
                        logger.error("Max retry attempts reached, aborting.")
                        raise
        return wrapper
    return decorator