python-backoff 2.3.1__py3-none-any.whl → 2.4.0__py3-none-any.whl

backoff/_common.py

@@ -1,8 +1,29 @@
+from __future__ import annotations
+
 import functools
 import logging
-import sys
+import time
 import traceback
 import warnings
+from typing import TYPE_CHECKING, Any, Callable, TypeVar
+
+if TYPE_CHECKING:
+    import sys
+    from collections.abc import Generator, Iterable
+
+    from backoff._typing import (
+        ContextDetails,
+        Details,
+        _ContextHandler,
+        _Jitterer,
+        _MaybeCallable,
+        _WaitGenerator,
+    )
+
+    if sys.version_info >= (3, 11):
+        from typing import Self
+    else:
+        from typing_extensions import Self
 
 
 # Use module-specific logger with a default null handler.
@@ -10,26 +31,38 @@ _logger = logging.getLogger("backoff")
 _logger.addHandler(logging.NullHandler())  # pragma: no cover
 _logger.setLevel(logging.INFO)
 
+T = TypeVar("T")
+_HandlerT = TypeVar("_HandlerT")
+
 
 # Evaluate arg that can be either a fixed value or a callable.
-def _maybe_call(f, *args, **kwargs):
+def _maybe_call(f: _MaybeCallable[T] | None, *args: Any, **kwargs: Any) -> T | None:
     if callable(f):
         try:
-            return f(*args, **kwargs)
+            return f(*args, **kwargs)  # ty:ignore[call-top-callable, invalid-return-type]
         except TypeError:
-            return f
+            return f  # type: ignore[return-value] # ty:ignore[invalid-return-type]
     else:
         return f
 
 
-def _init_wait_gen(wait_gen, wait_gen_kwargs):
+def _init_wait_gen(
+    wait_gen: _WaitGenerator,
+    wait_gen_kwargs: dict[str, Any],
+) -> Generator[float, Any, None]:
     kwargs = {k: _maybe_call(v) for k, v in wait_gen_kwargs.items()}
     initialized = wait_gen(**kwargs)
     initialized.send(None)  # Initialize with an empty send
     return initialized
 
 
-def _next_wait(wait, send_value, jitter, elapsed, max_time):
+def _next_wait(
+    wait: Generator[float, None, None],
+    send_value: Any,
+    jitter: _Jitterer | None,
+    elapsed: float,
+    max_time: float | None,
+) -> float:
     value = wait.send(send_value)
     try:
         seconds = jitter(value) if jitter is not None else value
@@ -42,7 +75,7 @@ def _next_wait(wait, send_value, jitter, elapsed, max_time):
             stacklevel=2,
         )
 
-        seconds = value + jitter()
+        seconds = value + jitter()  # type: ignore[call-arg, misc] # ty:ignore[missing-argument, call-non-callable]
 
     # don't sleep longer than remaining allotted max_time
     if max_time is not None:
@@ -51,7 +84,103 @@ def _next_wait(wait, send_value, jitter, elapsed, max_time):
     return seconds
 
 
-def _prepare_logger(logger):
+class _RetryState:
+    """Bookkeeping shared by the sync and async retry loops.
+
+    Tracks try count and elapsed time, and owns the initialized wait
+    generator, so `retry_predicate`/`retry_exception` only need to drive
+    the call/check/sleep sequence around it.
+    """
+
+    __slots__ = (
+        "elapsed",
+        "max_time",
+        "max_tries",
+        "start",
+        "tries",
+        "wait",
+    )
+
+    def __init__(
+        self,
+        wait_gen: _WaitGenerator,
+        wait_gen_kwargs: dict[str, Any],
+        *,
+        max_tries: _MaybeCallable[int] | None = None,
+        max_time: _MaybeCallable[float] | None = None,
+    ) -> None:
+        self.tries = 0
+        self.elapsed: float = 0
+        self.start = time.monotonic()
+        self.max_tries = _maybe_call(max_tries)
+        self.max_time = _maybe_call(max_time)
+        self.wait = _init_wait_gen(wait_gen, wait_gen_kwargs)
+
+    def start_attempt(self) -> None:
+        self.tries += 1
+
+    def record_elapsed(self) -> float:
+        self.elapsed = time.monotonic() - self.start
+        return self.elapsed
+
+    def exhausted(self) -> bool:
+        max_tries_exceeded = self.tries == self.max_tries
+        max_time_exceeded = self.max_time is not None and self.elapsed >= self.max_time
+        return max_tries_exceeded or max_time_exceeded
+
+    def next_wait(self, send_value: Any, jitter: _Jitterer | None) -> float:
+        return _next_wait(self.wait, send_value, jitter, self.elapsed, self.max_time)
+
+
+class _Attempt:
+    """A single attempt yielded by `retry_context`/`aretry_context`.
+
+    Used as `with attempt: ...`. Exceptions matching `exception_types` are
+    caught here so the driving generator (not this object) decides whether
+    to retry, sleep, or let the exception propagate; anything else, or no
+    exception at all, is left for the `with` block's normal exit behavior.
+    """
+
+    __slots__ = (
+        "exception",
+        "exception_types",
+    )
+
+    def __init__(
+        self,
+        exception_types: type[Exception] | tuple[type[Exception], ...],
+    ) -> None:
+        self.exception_types = exception_types
+        self.exception: BaseException | None = None
+
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: object,
+    ) -> bool:
+        if exc_type is None:
+            # self.outcome = _Success()
+            return False
+
+        if not issubclass(exc_type, self.exception_types):
+            return False  # not ours; propagate immediately
+
+        self.exception = exc
+        return True  # suppress for now; the driving generator decides next
+
+
+def _dispatch_handlers(handlers: Iterable[_ContextHandler], **details: Any) -> None:
+    for hdlr in handlers:
+        hdlr(details)  # type: ignore[arg-type] # ty:ignore[invalid-argument-type]
+
+
+def _prepare_logger(
+    logger: str | logging.Logger | logging.LoggerAdapter | None,
+) -> logging.Logger | logging.LoggerAdapter | None:
     if isinstance(logger, str):
         logger = logging.getLogger(logger)
     return logger
@@ -60,22 +189,23 @@ def _prepare_logger(logger):
 # Configure handler list with user specified handler and optionally
 # with a default handler bound to the specified logger.
 def _config_handlers(
-    user_handlers,
+    user_handlers: _HandlerT | Iterable[_HandlerT] | None,
     *,
-    default_handler=None,
-    logger=None,
-    log_level=None,
-):
-    handlers = []
+    default_handler: Callable[..., None] | None = None,
+    logger: logging.Logger | logging.LoggerAdapter | None = None,
+    log_level: int | None = None,
+) -> list[_HandlerT]:
+    handlers: list[_HandlerT] = []
     if logger is not None:
         assert log_level is not None, "Log level is not specified"
+        assert default_handler is not None, "Default handler is not specified"
         # bind the specified logger to the default log handler
         log_handler = functools.partial(
             default_handler,
             logger=logger,
             log_level=log_level,
         )
-        handlers.append(log_handler)
+        handlers.append(log_handler)  # type: ignore[arg-type] # ty:ignore[invalid-argument-type]
 
     if user_handlers is None:
         return handlers
@@ -84,22 +214,28 @@ def _config_handlers(
     # or a single handler. either way append them to the list.
     if hasattr(user_handlers, "__iter__"):
         # add all handlers in the iterable
-        handlers += list(user_handlers)
+        # pyrefly: ignore [bad-argument-type]
+        handlers += list(user_handlers)  # ty:ignore[invalid-argument-type]
     else:
         # append a single handler
+        # pyrefly: ignore [bad-argument-type]
         handlers.append(user_handlers)
 
     return handlers
 
 
 # Default backoff handler
-def _log_backoff(details, logger, log_level):
+def _log_backoff(
+    details: Details,
+    logger: logging.Logger | logging.LoggerAdapter,
+    log_level: int,
+) -> None:
     msg = "Backing off %s(...) for %.1fs (%s)"
-    log_args = [details["target"].__name__, details["wait"]]
+    log_args = [details["target"].__name__, details["wait"]]  # ty:ignore[unresolved-attribute]
 
-    exc_typ, exc, _ = sys.exc_info()
+    exc = details.get("exception")
     if exc is not None:
-        exc_fmt = traceback.format_exception_only(exc_typ, exc)[-1]
+        exc_fmt = traceback.format_exception_only(type(exc), exc)[-1]
         log_args.append(exc_fmt.rstrip("\n"))
     else:
         log_args.append(details["value"])
@@ -107,15 +243,49 @@ def _log_backoff(details, logger, log_level):
 
 
 # Default giveup handler
-def _log_giveup(details, logger, log_level):
+def _log_giveup(
+    details: Details,
+    logger: logging.Logger | logging.LoggerAdapter,
+    log_level: int,
+) -> None:
     msg = "Giving up %s(...) after %d tries (%s)"
-    log_args = [details["target"].__name__, details["tries"]]
+    log_args = [details["target"].__name__, details["tries"]]  # ty:ignore[unresolved-attribute]
 
-    exc_typ, exc, _ = sys.exc_info()
+    exc = details.get("exception")
     if exc is not None:
-        exc_fmt = traceback.format_exception_only(exc_typ, exc)[-1]
+        exc_fmt = traceback.format_exception_only(type(exc), exc)[-1]
         log_args.append(exc_fmt.rstrip("\n"))
     else:
         log_args.append(details["value"])
 
     logger.log(log_level, msg, *log_args)
+
+
+# Default backoff handler for retry_context/aretry_context (no wrapped
+# callable, so no name to log; the exception is read from `details`
+# directly since it's no longer the active exception by this point).
+def _log_backoff_context(
+    details: ContextDetails,
+    logger: logging.Logger | logging.LoggerAdapter,
+    log_level: int,
+) -> None:
+    logger.log(
+        log_level,
+        "Backing off retry_context(...) for %.1fs (%s)",
+        details["wait"],
+        details.get("exception"),
+    )
+
+
+# Default giveup handler for retry_context/aretry_context.
+def _log_giveup_context(
+    details: ContextDetails,
+    logger: logging.Logger | logging.LoggerAdapter,
+    log_level: int,
+) -> None:
+    logger.log(
+        log_level,
+        "Giving up retry_context(...) after %d tries (%s)",
+        details["tries"],
+        details.get("exception"),
+    )

backoff/_decorator.py

@@ -1,43 +1,62 @@
+from __future__ import annotations
+
 import inspect
 import logging
 import operator
-from typing import Any, Callable, Iterable, Optional, Type, Union
+from typing import TYPE_CHECKING, Any, Callable, TypeVar
 
+from backoff import _async, _sync
 from backoff._common import (
-    _prepare_logger,
     _config_handlers,
     _log_backoff,
+    _log_backoff_context,
     _log_giveup,
+    _log_giveup_context,
+    _prepare_logger,
 )
 from backoff._jitter import full_jitter
-from backoff import _async, _sync
-from backoff._typing import (
-    _CallableT,
-    _Handler,
-    _Jitterer,
-    _MaybeCallable,
-    _MaybeLogger,
-    _MaybeSequence,
-    _Predicate,
-    _WaitGenerator,
-)
+from backoff._wait_gen import expo
+
+if TYPE_CHECKING:
+    import sys
+    from collections.abc import AsyncGenerator, Generator, Iterable
+
+    from backoff._common import _Attempt
+    from backoff._typing import (
+        _ContextHandler,
+        _Handler,
+        _Jitterer,
+        _MaybeCallable,
+        _MaybeLogger,
+        _MaybeTuple,
+        _Predicate,
+        _WaitGenerator,
+    )
+
+    if sys.version_info >= (3, 10):
+        from typing import ParamSpec
+    else:
+        from typing_extensions import ParamSpec
+
+    T = TypeVar("T")
+    P = ParamSpec("P")
 
 
 def on_predicate(
     wait_gen: _WaitGenerator,
     predicate: _Predicate[Any] = operator.not_,
     *,
-    max_tries: Optional[_MaybeCallable[int]] = None,
-    max_time: Optional[_MaybeCallable[float]] = None,
-    jitter: Union[_Jitterer, None] = full_jitter,
-    on_success: Union[_Handler, Iterable[_Handler], None] = None,
-    on_backoff: Union[_Handler, Iterable[_Handler], None] = None,
-    on_giveup: Union[_Handler, Iterable[_Handler], None] = None,
+    max_tries: _MaybeCallable[int] | None = None,
+    max_time: _MaybeCallable[float] | None = None,
+    jitter: _Jitterer | None = full_jitter,
+    on_success: _Handler | Iterable[_Handler] | None = None,
+    on_backoff: _Handler | Iterable[_Handler] | None = None,
+    on_giveup: _Handler | Iterable[_Handler] | None = None,
     logger: _MaybeLogger = "backoff",
     backoff_log_level: int = logging.INFO,
     giveup_log_level: int = logging.ERROR,
     **wait_gen_kwargs: Any,
-) -> Callable[[_CallableT], _CallableT]:
+) -> Callable[[Callable[P, T]], Callable[P, T]]:
     """Returns decorator for backoff and retry triggered by predicate.
 
     Args:
@@ -82,7 +101,7 @@ def on_predicate(
             This is useful for runtime configuration.
     """
 
-    def decorate(target):
+    def decorate(target: Callable[P, T]) -> Callable[P, T]:
         nonlocal logger, on_success, on_backoff, on_giveup
 
         logger = _prepare_logger(logger)
@@ -124,21 +143,21 @@ def on_predicate(
 
 def on_exception(
     wait_gen: _WaitGenerator,
-    exception: _MaybeSequence[Type[Exception]],
+    exception: _MaybeTuple[type[Exception]],
     *,
-    max_tries: Optional[_MaybeCallable[int]] = None,
-    max_time: Optional[_MaybeCallable[float]] = None,
-    jitter: Union[_Jitterer, None] = full_jitter,
+    max_tries: _MaybeCallable[int] | None = None,
+    max_time: _MaybeCallable[float] | None = None,
+    jitter: _Jitterer | None = full_jitter,
     giveup: _Predicate[Exception] = lambda e: False,
-    on_success: Union[_Handler, Iterable[_Handler], None] = None,
-    on_backoff: Union[_Handler, Iterable[_Handler], None] = None,
-    on_giveup: Union[_Handler, Iterable[_Handler], None] = None,
+    on_success: _Handler | Iterable[_Handler] | None = None,
+    on_backoff: _Handler | Iterable[_Handler] | None = None,
+    on_giveup: _Handler | Iterable[_Handler] | None = None,
     raise_on_giveup: bool = True,
     logger: _MaybeLogger = "backoff",
     backoff_log_level: int = logging.INFO,
     giveup_log_level: int = logging.ERROR,
     **wait_gen_kwargs: Any,
-) -> Callable[[_CallableT], _CallableT]:
+) -> Callable[[Callable[P, T]], Callable[P, T]]:
     """Returns decorator for backoff and retry triggered by exception.
 
     Args:
@@ -185,7 +204,7 @@ def on_exception(
             This is useful for runtime configuration.
     """
 
-    def decorate(target):
+    def decorate(target: Callable[P, T]) -> Callable[P, T]:
         nonlocal logger, on_success, on_backoff, on_giveup
 
         logger = _prepare_logger(logger)
@@ -225,3 +244,155 @@ def on_exception(
 
     # Return a function which decorates a target with a retry loop.
     return decorate
+
+
+def retry_context(
+    exception: _MaybeTuple[type[Exception]] = Exception,
+    wait_gen: _WaitGenerator = expo,
+    *,
+    max_tries: _MaybeCallable[int] | None = None,
+    max_time: _MaybeCallable[float] | None = None,
+    jitter: _Jitterer | None = full_jitter,
+    giveup: _Predicate[BaseException] = lambda e: False,
+    on_success: _ContextHandler | Iterable[_ContextHandler] | None = None,
+    on_backoff: _ContextHandler | Iterable[_ContextHandler] | None = None,
+    on_giveup: _ContextHandler | Iterable[_ContextHandler] | None = None,
+    raise_on_giveup: bool = True,
+    logger: _MaybeLogger = "backoff",
+    backoff_log_level: int = logging.INFO,
+    giveup_log_level: int = logging.ERROR,
+    **wait_gen_kwargs: Any,
+) -> Generator[_Attempt, None, None]:
+    """Returns a generator of retry attempts, for direct use with a `for` loop.
+
+    Unlike `on_exception`, this doesn't wrap a whole function; it lets a
+    caller retry an arbitrary block of code:
+
+        for attempt in backoff.retry_context(ValueError, backoff.expo):
+            with attempt:
+                do_something()
+
+    Each `attempt` is a context manager: exceptions matching `exception`
+    are caught, and the loop either sleeps and retries or lets the
+    exception (or a different one) propagate once retries are exhausted.
+    Succeeding (no exception raised in the `with` block) ends the loop.
+
+    Args:
+        exception: An exception type (or tuple of types) which triggers
+            backoff.
+        wait_gen: A generator yielding successive wait times in seconds.
+        max_tries: The maximum number of attempts to make before giving
+            up. Once exhausted, the exception will be allowed to escape.
+            The default value of None means there is no limit to the
+            number of tries. If a callable is passed, it will be
+            evaluated at runtime and its return value used.
+        max_time: The maximum total amount of time to try for before
+            giving up. Once expired, the exception will be allowed to
+            escape. If a callable is passed, it will be evaluated at
+            runtime and its return value used.
+        jitter: A function of the value yielded by wait_gen returning
+            the actual time to wait. Jittered by default using
+            full_jitter; disable with jitter=None.
+        giveup: Function accepting an exception instance and returning
+            whether or not to give up. Optional. The default is to
+            always continue.
+        on_success: Callable (or iterable of callables) with a unary
+            signature called on success. The parameter is a dict with
+            `tries` and `elapsed`.
+        on_backoff: Callable (or iterable of callables) called on
+            backoff. The parameter dict additionally has `wait` and
+            `exception`.
+        on_giveup: Callable (or iterable of callables) called when
+            giving up. The parameter dict additionally has `exception`.
+        raise_on_giveup: Boolean indicating whether the registered
+            exception should be raised on giveup. Defaults to `True`.
+        logger: Name or Logger object to log to. Defaults to 'backoff'.
+        backoff_log_level: log level for the backoff event. Defaults to "INFO"
+        giveup_log_level: log level for the give up event. Defaults to "ERROR"
+        **wait_gen_kwargs: Any additional keyword args specified will be
+            passed to wait_gen when it is initialized.
+    """
+    logger = _prepare_logger(logger)
+    on_success = _config_handlers(on_success)
+    on_backoff = _config_handlers(
+        on_backoff,
+        default_handler=_log_backoff_context,
+        logger=logger,
+        log_level=backoff_log_level,
+    )
+    on_giveup = _config_handlers(
+        on_giveup,
+        default_handler=_log_giveup_context,
+        logger=logger,
+        log_level=giveup_log_level,
+    )
+
+    return _sync.retry_context(
+        exception,
+        wait_gen,
+        max_tries=max_tries,
+        max_time=max_time,
+        jitter=jitter,
+        giveup=giveup,
+        on_success=on_success,
+        on_backoff=on_backoff,
+        on_giveup=on_giveup,
+        raise_on_giveup=raise_on_giveup,
+        wait_gen_kwargs=wait_gen_kwargs,
+    )
+
+
+def aretry_context(
+    exception: _MaybeTuple[type[Exception]] = Exception,
+    wait_gen: _WaitGenerator = expo,
+    *,
+    max_tries: _MaybeCallable[int] | None = None,
+    max_time: _MaybeCallable[float] | None = None,
+    jitter: _Jitterer | None = full_jitter,
+    giveup: _Predicate[BaseException] = lambda e: False,
+    on_success: _ContextHandler | Iterable[_ContextHandler] | None = None,
+    on_backoff: _ContextHandler | Iterable[_ContextHandler] | None = None,
+    on_giveup: _ContextHandler | Iterable[_ContextHandler] | None = None,
+    raise_on_giveup: bool = True,
+    logger: _MaybeLogger = "backoff",
+    backoff_log_level: int = logging.INFO,
+    giveup_log_level: int = logging.ERROR,
+    **wait_gen_kwargs: Any,
+) -> AsyncGenerator[_Attempt, None]:
+    """Async counterpart to `retry_context`, for use with `async for`.
+
+        async for attempt in backoff.aretry_context(ValueError, backoff.expo):
+            with attempt:
+                await do_something()
+
+    `giveup` and the handlers may be sync or async callables. See
+    `retry_context` for the full argument reference.
+    """
+    logger = _prepare_logger(logger)
+    on_success = _config_handlers(on_success)
+    on_backoff = _config_handlers(
+        on_backoff,
+        default_handler=_log_backoff_context,
+        logger=logger,
+        log_level=backoff_log_level,
+    )
+    on_giveup = _config_handlers(
+        on_giveup,
+        default_handler=_log_giveup_context,
+        logger=logger,
+        log_level=giveup_log_level,
+    )
+
+    return _async.aretry_context(
+        exception,
+        wait_gen,
+        max_tries=max_tries,
+        max_time=max_time,
+        jitter=jitter,
+        giveup=giveup,
+        on_success=on_success,
+        on_backoff=on_backoff,
+        on_giveup=on_giveup,
+        raise_on_giveup=raise_on_giveup,
+        wait_gen_kwargs=wait_gen_kwargs,
+    )