Add retry_context (#12)

* Add retry_context

* Add changelog entry

* Make context retries respect config

* Cache no-op stop

* Add narrative docs

* Fix typo

Author: hynek

CHANGELOG.md

@@ -19,6 +19,8 @@ The **third number** is for emergencies when we need to start branches for older
   [#9](https://github.com/hynek/stamina/pull/9)
 - Async support.
   [#10](https://github.com/hynek/stamina/pull/10)
+- Retries of arbitrary blocks using (async) `for` loops and context managers.
+  [#12](https://github.com/hynek/stamina/pull/12)
 
 
 ## [22.2.0](https://github.com/hynek/stamina/compare/22.1.0...22.2.0) - 2022-10-06

README.md

@@ -25,7 +25,7 @@ In practice, only a few knobs are needed (repeatedly!), though.
 
 ## Usage
 
-The API consists mainly of the `stamina.retry()` decorator:
+The API consists mainly of the `stamina.retry()` decorator for retrying functions and methods, and the `stamina.retry_context()` iterator / context manager combo for retrying arbitrary code blocks:
 
 <!-- example-start -->
 ```python
@@ -42,6 +42,11 @@ def do_it(code: int) -> httpx.Response:
 
 # reveal_type(do_it)
 # note: Revealed type is "def (code: builtins.int) -> httpx._models.Response"
+
+for attempt in stamina.retry_context(on=httpx.HTTPError, attempts=3):
+    with attempt:
+        resp = httpx.get(f"https://httpbin.org/status/{code}")
+        resp.raise_for_status()
 ```
 
 Async works the same way:
@@ -57,10 +62,16 @@ async def do_it_async(code: int) -> httpx.Response:
 
 # reveal_type(do_it_async)
 # note: Revealed type is "def (code: builtins.int) -> typing.Coroutine[Any, Any, httpx._models.Response]"
+
+async for attempt in stamina.retry_context(on=httpx.HTTPError, attempts=3):
+    with attempt:
+        async with httpx.AsyncClient() as client:
+            resp = await client.get(f"https://httpbin.org/status/{code}")
+        resp.raise_for_status()
 ```
 <!-- example-end -->
 
-The decorator takes the following arguments (**all time-based arguments are floats of seconds**):
+Both `retry()` and `retry_context()` take the following arguments (**all time-based arguments are floats of seconds**):
 
 **on**: An Exception or a tuple of Exceptions on which the decorated callable will be retried.
 There is no default – you _must_ pass this explicitly.

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 
 [tool.supported-pythons]
-# These are canonical truth that is then used by Cog to generate metadata & CI
+# This is the canonical truth that is then used by Cog to generate metadata & CI
 # config.
 min = "3.8"
 max = "3.12"

src/stamina/__init__.py

@@ -3,8 +3,14 @@
 # SPDX-License-Identifier: MIT
 
 from ._config import is_active, set_active
-from ._core import retry
+from ._core import retry, retry_context
 from ._instrumentation import RETRY_COUNTER
 
 
-__all__ = ["retry", "is_active", "set_active", "RETRY_COUNTER"]
+__all__ = [
+    "retry",
+    "retry_context",
+    "is_active",
+    "set_active",
+    "RETRY_COUNTER",
+]

src/stamina/_core.py

@@ -7,6 +7,7 @@
 import sys
 
 from collections.abc import Callable
+from dataclasses import dataclass, replace
 from functools import wraps
 from inspect import iscoroutinefunction
 from typing import Iterable, TypeVar
@@ -32,6 +33,164 @@
 P = ParamSpec("P")
 
 
+def retry_context(
+    on: type[Exception] | tuple[type[Exception], ...],
+    attempts: int | None = 10,
+    timeout: float | int | None = 45.0,
+    wait_initial: float | int = 0.1,
+    wait_max: float | int = 5.0,
+    wait_jitter: float | int = 1.0,
+    wait_exp_base: float | int = 2.0,
+) -> _RetryContextIterator:
+    """
+    Iterator that yields context managers that can be used to retry code
+    blocks.
+
+    Arguments have the same meaning as for :func:`retry`.
+
+    .. versionadded:: 23.1.0
+    """
+
+    return _RetryContextIterator.from_params(
+        on=on,
+        attempts=attempts,
+        timeout=timeout,
+        wait_initial=wait_initial,
+        wait_max=wait_max,
+        wait_jitter=wait_jitter,
+        wait_exp_base=wait_exp_base,
+        name="<context block>",
+        args=(),
+        kw={},
+    )
+
+
+@dataclass
+class _RetryContextIterator:
+    __slots__ = ("_tenacity_kw", "_name", "_args", "_kw")
+    _tenacity_kw: dict[str, object]
+    _name: str
+    _args: tuple[object, ...]
+    _kw: dict[str, object]
+
+    @classmethod
+    def from_params(
+        cls,
+        on: type[Exception] | tuple[type[Exception], ...],
+        attempts: int | None,
+        timeout: float | int | None,
+        wait_initial: float | int,
+        wait_max: float | int,
+        wait_jitter: float | int,
+        wait_exp_base: float | int,
+        name: str,
+        args: tuple[object, ...],
+        kw: dict[str, object],
+    ) -> _RetryContextIterator:
+        return cls(
+            _name=name,
+            _args=args,
+            _kw=kw,
+            _tenacity_kw={
+                "retry": _t.retry_if_exception_type(on),
+                "wait": _t.wait_exponential_jitter(
+                    initial=wait_initial,
+                    max=wait_max,
+                    exp_base=wait_exp_base,
+                    jitter=wait_jitter,
+                ),
+                "stop": _make_stop(attempts=attempts, timeout=timeout),
+                "reraise": True,
+            },
+        )
+
+    _STOP_NO_RETRY = _t.stop_after_attempt(1)
+
+    def __iter__(self) -> _t.Retrying:
+        if not _CONFIG.is_active:
+            return _t.Retrying(
+                reraise=True, stop=self._STOP_NO_RETRY
+            ).__iter__()
+
+        return _t.Retrying(
+            before_sleep=_make_before_sleep(
+                self._name, _CONFIG.on_retry, self._args, self._kw
+            )
+            if _CONFIG.on_retry
+            else None,
+            **self._tenacity_kw,
+        ).__iter__()
+
+    def __aiter__(self) -> _t.AsyncRetrying:
+        if not _CONFIG.is_active:
+            return _t.AsyncRetrying(
+                reraise=True, stop=self._STOP_NO_RETRY
+            ).__aiter__()
+
+        return _t.AsyncRetrying(
+            before_sleep=_make_before_sleep(
+                self._name, _CONFIG.on_retry, self._args, self._kw
+            )
+            if _CONFIG.on_retry
+            else None,
+            **self._tenacity_kw,
+        ).__aiter__()
+
+    def with_name(
+        self, name: str, args: tuple[object, ...], kw: dict[str, object]
+    ) -> _RetryContextIterator:
+        """
+        Recreate ourselves with a new name and arguments.
+        """
+        return replace(self, _name=name, _args=args, _kw=kw)
+
+
+def _make_before_sleep(
+    name: str, on_retry: Iterable[RetryHook], args: object, kw: object
+) -> Callable[[_t.RetryCallState], None]:
+    """
+    Create a `before_sleep` callback function that runs our `RetryHook`s with
+    the necessary arguments.
+    """
+
+    def before_sleep(rcs: _t.RetryCallState) -> None:
+        attempt = rcs.attempt_number
+        exc = rcs.outcome.exception()
+        backoff = rcs.idle_for
+
+        for hook in on_retry:
+            hook(
+                attempt=attempt,
+                backoff=backoff,
+                exc=exc,
+                name=name,
+                args=args,
+                kwargs=kw,
+            )
+
+    return before_sleep
+
+
+def _make_stop(*, attempts: int | None, timeout: float | None) -> _t.stop_base:
+    """
+    Combine *attempts* and *timeout* into one stop condition.
+    """
+    stops = []
+
+    if attempts:
+        stops.append(_t.stop_after_attempt(attempts))
+    if timeout:
+        stops.append(_t.stop_after_delay(timeout))
+
+    if len(stops) > 1:
+        return _t.stop_any(*stops)
+
+    if not stops:
+        return _t.stop_never
+
+    return stops[0]
+
+
 def retry(
     *,
     on: type[Exception] | tuple[type[Exception], ...],
@@ -62,14 +221,18 @@ def do_it(code: int) -> httpx.Response:
 
                 return resp
     """
-    retry_ = _t.retry_if_exception_type(on)
-    wait = _t.wait_exponential_jitter(
-        initial=wait_initial,
-        max=wait_max,
-        exp_base=wait_exp_base,
-        jitter=wait_jitter,
+    retry_ctx = _RetryContextIterator.from_params(
+        on=on,
+        attempts=attempts,
+        timeout=timeout,
+        wait_initial=wait_initial,
+        wait_max=wait_max,
+        wait_jitter=wait_jitter,
+        wait_exp_base=wait_exp_base,
+        name="<unknown>",
+        args=(),
+        kw={},
     )
-    stop = _make_stop(attempts=attempts, timeout=timeout)
 
     def retry_decorator(wrapped: Callable[P, T]) -> Callable[P, T]:
         name = guess_name(wrapped)
@@ -78,19 +241,8 @@ def retry_decorator(wrapped: Callable[P, T]) -> Callable[P, T]:
 
             @wraps(wrapped)
             def sync_inner(*args: P.args, **kw: P.kwargs) -> T:  # type: ignore[return]
-                if not _CONFIG.is_active:
-                    return wrapped(*args, **kw)
-
-                for attempt in _t.Retrying(  # noqa: RET503
-                    retry=retry_,
-                    wait=wait,
-                    stop=stop,
-                    reraise=True,
-                    before_sleep=_make_before_sleep(
-                        name, _CONFIG.on_retry, args, kw
-                    )
-                    if _CONFIG.on_retry
-                    else None,
+                for attempt in retry_ctx.with_name(  # noqa: RET503
+                    name, args, kw
                 ):
                     with attempt:
                         return wrapped(*args, **kw)
@@ -99,69 +251,12 @@ def sync_inner(*args: P.args, **kw: P.kwargs) -> T:  # type: ignore[return]
 
         @wraps(wrapped)
         async def async_inner(*args: P.args, **kw: P.kwargs) -> T:  # type: ignore[return]
-            if not _CONFIG.is_active:
-                return await wrapped(*args, **kw)  # type: ignore[no-any-return,misc]
-
-            async for attempt in _t.AsyncRetrying(  # noqa: RET503
-                retry=retry_,
-                wait=wait,
-                stop=stop,
-                reraise=True,
-                before_sleep=_make_before_sleep(
-                    name, _CONFIG.on_retry, args, kw
-                )
-                if _CONFIG.on_retry
-                else None,
+            async for attempt in retry_ctx.with_name(  # noqa: RET503
+                name, args, kw
             ):
                 with attempt:
                     return await wrapped(*args, **kw)  # type: ignore[misc,no-any-return]
 
         return async_inner  # type: ignore[return-value]
 
     return retry_decorator
-
-
-def _make_stop(*, attempts: int | None, timeout: float | None) -> _t.stop_base:
-    """
-    Combine *attempts* and *timeout* into one stop condition.
-    """
-    stops = []
-
-    if attempts:
-        stops.append(_t.stop_after_attempt(attempts))
-    if timeout:
-        stops.append(_t.stop_after_delay(timeout))
-
-    if len(stops) > 1:
-        return _t.stop_any(*stops)
-
-    if not stops:
-        return _t.stop_never
-
-    return stops[0]
-
-
-def _make_before_sleep(
-    name: str, on_retry: Iterable[RetryHook], args: object, kw: object
-) -> Callable[[_t.RetryCallState], None]:
-    """
-    Create a `before_sleep` callback function that runs our `RetryHook`s with
-    the necessary arguments.
-    """
-
-    def before_sleep(rcs: _t.RetryCallState) -> None:
-        attempt = rcs.attempt_number
-        exc = rcs.outcome.exception()
-        backoff = rcs.idle_for
-
-        for hook in on_retry:
-            hook(
-                attempt=attempt,
-                backoff=backoff,
-                exc=exc,
-                name=name,
-                args=args,
-                kwargs=kw,
-            )
-
-    return before_sleep