Allow instrumentation hooks to be contextmanagers

fixes #84

Author: hynek

CHANGELOG.md

@@ -25,6 +25,9 @@ You can find our backwards-compatibility policy [here](https://github.com/hynek/
 - `stamina.set_testing()` can now be used as a context manager.
   [#94](https://github.com/hynek/stamina/pull/94)
 
+- Instrumentation hooks can now can return context managers.
+  If they do, they are entered before going to sleep and exited after waking up.
+
 
 ## [24.3.0](https://github.com/hynek/stamina/compare/24.2.0...24.3.0) - 2024-08-27
 

pyproject.toml

@@ -27,7 +27,7 @@ classifiers = [
 dependencies = ["tenacity", "typing-extensions; python_version < '3.10'"]
 
 [project.optional-dependencies]
-tests = ["pytest", "anyio"]
+tests = ["pytest", "anyio", "dirty-equals"]
 typing = ["mypy >= 1.4"]
 docs = [
   "sphinx >= 7.2.2",

src/stamina/_core.py

@@ -4,10 +4,12 @@
 
 from __future__ import annotations
 
+import contextlib
 import datetime as dt
 import random
 import sys
 
+from contextlib import AbstractContextManager
 from dataclasses import dataclass, replace
 from functools import wraps
 from inspect import iscoroutinefunction
@@ -400,6 +402,7 @@ class _RetryContextIterator:
     __slots__ = (
         "_args",
         "_attempts",
+        "_cms_to_exit",
         "_kw",
         "_name",
         "_t_a_retrying",
@@ -421,6 +424,8 @@ class _RetryContextIterator:
     _wait_max: float
     _wait_exp_base: float
 
+    _cms_to_exit: list[AbstractContextManager[None]]
+
     @classmethod
     def from_params(
         cls,
@@ -473,6 +478,7 @@ def from_params(
                 "reraise": True,
             },
             _t_a_retrying=_LAZY_NO_ASYNC_RETRY,
+            _cms_to_exit=[],
         )
 
         inst._t_kw["wait"] = inst._jittered_backoff_for_rcs
@@ -501,17 +507,24 @@ def _apply_maybe_test_mode_to_tenacity_kw(
 
         return t_kw
 
+    def _exit_cms(self, _: _t.RetryCallState | None) -> None:
+        for cm in reversed(self._cms_to_exit):
+            cm.__exit__(None, None, None)
+
     def __iter__(self) -> Iterator[Attempt]:
         if not CONFIG.is_active:
             for r in _t.Retrying(reraise=True, stop=_STOP_NO_RETRY):
                 yield Attempt(r, None)
 
             return
 
+        before_sleep = _make_before_sleep(
+            self._name, CONFIG, self._args, self._kw, self._cms_to_exit
+        )
+
         for r in _t.Retrying(
-            before_sleep=_make_before_sleep(
-                self._name, CONFIG, self._args, self._kw
-            ),
+            before=self._exit_cms,
+            before_sleep=before_sleep,
             **self._apply_maybe_test_mode_to_tenacity_kw(CONFIG.testing),
         ):
             yield Attempt(r, self._backoff_for_attempt_number)
@@ -520,8 +533,9 @@ def __aiter__(self) -> AsyncIterator[Attempt]:
         if CONFIG.is_active:
             self._t_a_retrying = _t.AsyncRetrying(
                 sleep=_smart_sleep,
+                before=self._exit_cms,
                 before_sleep=_make_before_sleep(
-                    self._name, CONFIG, self._args, self._kw
+                    self._name, CONFIG, self._args, self._kw, self._cms_to_exit
                 ),
                 **self._apply_maybe_test_mode_to_tenacity_kw(CONFIG.testing),
             )
@@ -583,10 +597,15 @@ def _make_before_sleep(
     config: _Config,
     args: tuple[object, ...],
     kw: dict[str, object],
+    hook_cms: list[contextlib.AbstractContextManager[None]],
 ) -> Callable[[_t.RetryCallState], None]:
     """
     Create a `before_sleep` callback function that runs our `RetryHook`s with
     the necessary arguments.
+
+    If a hook returns a context manager, it's entered before retries start and
+    exited after they finish by keeping track of the context managers in
+    *hook_cms*.
     """
 
     last_idle_for = 0.0
@@ -607,7 +626,11 @@ def before_sleep(rcs: _t.RetryCallState) -> None:
         )
 
         for hook in config.on_retry:
-            hook(details)
+            maybe_cm = hook(details)
+
+            if isinstance(maybe_cm, AbstractContextManager):
+                maybe_cm.__enter__()
+                hook_cms.append(maybe_cm)
 
         last_idle_for = rcs.idle_for
 

src/stamina/instrumentation/_data.py

@@ -4,6 +4,7 @@
 
 from __future__ import annotations
 
+from contextlib import AbstractContextManager
 from dataclasses import dataclass
 from typing import Callable, Protocol
 
@@ -75,7 +76,9 @@ class RetryHook(Protocol):
     .. versionadded:: 23.2.0
     """
 
-    def __call__(self, details: RetryDetails) -> None: ...
+    def __call__(
+        self, details: RetryDetails
+    ) -> None | AbstractContextManager[None]: ...
 
 
 @dataclass(frozen=True)

tests/test_instrumentation.py

@@ -4,16 +4,20 @@
 
 from __future__ import annotations
 
+from contextlib import contextmanager
+
 import pytest
 
+from dirty_equals import IsInstance
+
 import stamina
 
 from stamina.instrumentation import (
     RetryHookFactory,
     get_on_retry_hooks,
     set_on_retry_hooks,
 )
-from stamina.instrumentation._data import guess_name
+from stamina.instrumentation._data import RetryDetails, guess_name
 from stamina.instrumentation._hooks import get_default_hooks
 from stamina.instrumentation._structlog import init_structlog
 
@@ -191,3 +195,71 @@ def init():
             hook,
             delayed_hook,
         ) == get_on_retry_hooks()
+
+    def test_context_manager_hooks(self):
+        """
+        If a hook is a context manager, it's entered before retries start and exited
+        after they finish.
+        """
+        entered = False
+        exited = False
+        deets = []
+
+        @contextmanager
+        def cm1(details):
+            nonlocal entered
+            entered = True
+
+            deets.append(details)
+            yield
+
+            nonlocal exited
+            exited = True
+
+        class CM:
+            def __init__(self):
+                self.entered = False
+                self.exited = False
+                self.deets = []
+
+            def __call__(self, details):
+                self.deets.append(details)
+                return self
+
+            def __enter__(self):
+                self.entered = True
+
+            def __exit__(self, *_):
+                self.exited = True
+
+        cm2 = CM()
+
+        set_on_retry_hooks([cm1, cm2])
+
+        @stamina.retry(on=ValueError, wait_max=0, attempts=2)
+        def f():
+            raise ValueError
+
+        with pytest.raises(ValueError):
+            f()
+
+        assert entered
+        assert exited
+        assert cm2.entered
+        assert cm2.exited
+
+        assert (
+            [
+                RetryDetails(
+                    name="tests.test_instrumentation.TestSetOnRetryHooks.test_context_manager_hooks.<locals>.f",
+                    args=(),
+                    kwargs={},
+                    retry_num=1,
+                    wait_for=0.0,
+                    waited_so_far=0.0,
+                    caused_by=IsInstance(ValueError),
+                )
+            ]
+            == cm2.deets
+            == deets
+        )

tests/typing/api.py

@@ -11,6 +11,9 @@
 
 import datetime as dt
 
+from collections.abc import Generator
+from contextlib import contextmanager
+
 from stamina import (
     AsyncRetryingCaller,
     BoundAsyncRetryingCaller,
@@ -109,11 +112,16 @@ def hook(details: RetryDetails) -> None:
     return None
 
 
+@contextmanager
+def cm_hook(details: RetryDetails) -> Generator[None]:
+    yield
+
+
 def init() -> RetryHook:
     return hook
 
 
-set_on_retry_hooks([hook, RetryHookFactory(init)])
+set_on_retry_hooks([hook, RetryHookFactory(init), cm_hook])
 
 hooks: tuple[RetryHook, ...] = get_on_retry_hooks()