Hide Tenacity's AttemptManager from the users (#22)

* Hide Tenacity's AttemptManager from the users

* Add PR link

* Instantiate no-retry AsyncRetrying lazily

* Save some lookups

* Fix type

* Clarify

* This fits now

* Better wording

* Expose the number of the current attempt

* Stop vendoring

Co-authored-by: Tin Tvrtković <[email protected]>

---------

Co-authored-by: Tin Tvrtković <[email protected]>

Author: hynek

CHANGELOG.md

@@ -15,6 +15,13 @@ You can find our backwards-compatibility policy [here](https://github.com/hynek/
 
 ## [Unreleased](https://github.com/hynek/stamina/compare/23.1.0...HEAD)
 
+### Changed
+
+- Tenacity's internal `AttemptManager` object is no longer exposed to the user.
+  This was an oversight and never documented.
+  `stamina.retry_context()` now yields instances of `stamina.Attempt`.
+  [#22](https://github.com/hynek/stamina/pull/22)
+
 
 ## [23.1.0](https://github.com/hynek/stamina/compare/22.2.0...23.1.0) - 2023-07-04
 

docs/api.rst

@@ -5,6 +5,8 @@ API Reference
 
 .. autofunction:: retry
 .. autofunction:: retry_context
+.. autoclass:: Attempt
+   :members: num
 
 
 Configuration

src/stamina/__init__.py

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

src/stamina/_core.py

@@ -11,12 +11,8 @@
 from dataclasses import dataclass, replace
 from functools import wraps
 from inspect import iscoroutinefunction
-from typing import (
-    AsyncIterator,
-    Iterable,
-    Iterator,
-    TypeVar,
-)
+from types import TracebackType
+from typing import AsyncIterator, Iterable, Iterator, TypeVar
 
 import tenacity as _t
 
@@ -68,10 +64,59 @@ def retry_context(
     )
 
 
+@dataclass
+class Attempt:
+    """
+    A context manager that can be used to retry code blocks.
+
+    Instances are yielded by the :func:`stamina.retry_context` iterator.
+
+    .. versionadded:: 23.2.0
+    """
+
+    _t_attempt: _t.AttemptManager
+
+    @property
+    def num(self) -> int:
+        """
+        The number of the current attempt.
+        """
+        return self._t_attempt.retry_state.attempt_number  # type: ignore[no-any-return]
+
+    def __enter__(self) -> None:
+        return self._t_attempt.__enter__()  # type: ignore[no-any-return]
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> bool | None:
+        return self._t_attempt.__exit__(  # type: ignore[no-any-return]
+            exc_type, exc_value, traceback
+        )
+
+
+_STOP_NO_RETRY = _t.stop_after_attempt(1)
+
+
+class _LazyNoAsyncRetry:
+    """
+    Allows us a free null object pattern using non-retries and avoid None.
+    """
+
+    def __aiter__(self) -> _t.AsyncRetrying:
+        return _t.AsyncRetrying(reraise=True, stop=_STOP_NO_RETRY).__aiter__()
+
+
+_LAZY_NO_ASYNC_RETRY = _LazyNoAsyncRetry()
+
+
 @dataclass
 class _RetryContextIterator:
-    __slots__ = ("_tenacity_kw", "_name", "_args", "_kw")
-    _tenacity_kw: dict[str, object]
+    __slots__ = ("_t_kw", "_t_a_retrying", "_name", "_args", "_kw")
+    _t_kw: dict[str, object]
+    _t_a_retrying: _t.AsyncRetrying
     _name: str
     _args: tuple[object, ...]
     _kw: dict[str, object]
@@ -94,7 +139,7 @@ def from_params(
             _name=name,
             _args=args,
             _kw=kw,
-            _tenacity_kw={
+            _t_kw={
                 "retry": _t.retry_if_exception_type(on),
                 "wait": _t.wait_exponential_jitter(
                     initial=wait_initial.total_seconds()
@@ -116,45 +161,51 @@ def from_params(
                 ),
                 "reraise": True,
             },
+            _t_a_retrying=_LAZY_NO_ASYNC_RETRY,
         )
 
-    _STOP_NO_RETRY = _t.stop_after_attempt(1)
+    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 __iter__(self) -> Iterator[_t.AttemptManager]:
+    def __iter__(self) -> Iterator[Attempt]:
         if not _CONFIG.is_active:
-            yield from _t.Retrying(reraise=True, stop=self._STOP_NO_RETRY)
+            for r in _t.Retrying(
+                reraise=True, stop=_STOP_NO_RETRY
+            ):  # pragma: no cover -- it's always once + GeneratorExit
+                yield Attempt(r)
 
-        yield from _t.Retrying(
+        for r in _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,
-        )
+            **self._t_kw,
+        ):
+            yield Attempt(r)
+
+    def __aiter__(self) -> AsyncIterator[Attempt]:
+        if _CONFIG.is_active:
+            self._t_a_retrying = _t.AsyncRetrying(
+                before_sleep=_make_before_sleep(
+                    self._name, _CONFIG.on_retry, self._args, self._kw
+                )
+                if _CONFIG.on_retry
+                else None,
+                **self._t_kw,
+            )
 
-    def __aiter__(self) -> AsyncIterator[_t.AttemptManager]:
-        if not _CONFIG.is_active:
-            return _t.AsyncRetrying(  # type: ignore[no-any-return]
-                reraise=True, stop=self._STOP_NO_RETRY
-            ).__aiter__()
+        self._t_a_retrying = self._t_a_retrying.__aiter__()
 
-        return _t.AsyncRetrying(  # type: ignore[no-any-return]
-            before_sleep=_make_before_sleep(
-                self._name, _CONFIG.on_retry, self._args, self._kw
-            )
-            if _CONFIG.on_retry
-            else None,
-            **self._tenacity_kw,
-        ).__aiter__()
+        return self
 
-    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)
+    async def __anext__(self) -> Attempt:
+        return Attempt(await self._t_a_retrying.__anext__())
 
 
 def _make_before_sleep(

tests/test_async.py

@@ -124,6 +124,24 @@ async def f():
     assert 1 == num_called
 
 
+async def test_retry_inactive_ok():
+    """
+    If inactive, the happy path still works.
+    """
+    num_called = 0
+
+    @stamina.retry(on=Exception)
+    async def f():
+        nonlocal num_called
+        num_called += 1
+
+    stamina.set_active(False)
+
+    await f()
+
+    assert 1 == num_called
+
+
 async def test_retry_block():
     """
     Async retry_context blocks are retried.
@@ -133,6 +151,9 @@ async def test_retry_block():
     async for attempt in stamina.retry_context(on=ValueError, wait_max=0):
         with attempt:
             num_called += 1
+
+            assert num_called == attempt.num
+
             if num_called < 2:
                 raise ValueError
 

tests/test_sync.py

@@ -97,6 +97,24 @@ def f():
     assert 1 == num_called
 
 
+def test_retry_inactive_ok():
+    """
+    If inactive, the happy path still works.
+    """
+    num_called = 0
+
+    @stamina.retry(on=Exception)
+    def f():
+        nonlocal num_called
+        num_called += 1
+
+    stamina.set_active(False)
+
+    f()
+
+    assert 1 == num_called
+
+
 def test_retry_block():
     """
     Sync retry_context blocks are retried.
@@ -106,6 +124,9 @@ def test_retry_block():
     for attempt in stamina.retry_context(on=ValueError, wait_max=0):
         with attempt:
             i += 1
+
+            assert i == attempt.num
+
             if i < 2:
                 raise ValueError
 

tests/typing/api.py

@@ -98,7 +98,7 @@ def hook(
 
 for attempt in retry_context(on=ValueError, timeout=13):
     with attempt:
-        ...
+        x: int = attempt.num
 
 for attempt in retry_context(
     on=ValueError, timeout=dt.timedelta(seconds=13.0)
@@ -116,4 +116,4 @@ async def f() -> None:
         wait_jitter=one_sec,
     ):
         with attempt:
-            ...
+            pass