Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add support for async strategies and callbacks #433

Closed
wants to merge 4 commits into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
118 changes: 85 additions & 33 deletions tenacity/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,7 @@
if t.TYPE_CHECKING:
import types

from . import asyncio as tasyncio
from .retry import RetryBaseT
from .stop import StopBaseT
from .wait import WaitBaseT
Expand Down Expand Up @@ -279,6 +280,14 @@ def statistics(self) -> t.Dict[str, t.Any]:
self._local.statistics = t.cast(t.Dict[str, t.Any], {})
return self._local.statistics

@property
def iter_state(self) -> t.Dict[str, t.Any]:
try:
return self._local.iter_state # type: ignore[no-any-return]
except AttributeError:
self._local.iter_state = t.cast(t.Dict[str, t.Any], {})
return self._local.iter_state

def wraps(self, f: WrappedFn) -> WrappedFn:
"""Wrap a function for retrying.

Expand All @@ -303,45 +312,88 @@ def begin(self) -> None:
self.statistics["attempt_number"] = 1
self.statistics["idle_for"] = 0

def iter(self, retry_state: "RetryCallState") -> t.Union[DoAttempt, DoSleep, t.Any]: # noqa
fut = retry_state.outcome
if fut is None:
if self.before is not None:
self.before(retry_state)
return DoAttempt()

is_explicit_retry = fut.failed and isinstance(fut.exception(), TryAgain)
if not (is_explicit_retry or self.retry(retry_state)):
return fut.result()
def _add_action_func(self, fn: t.Callable[..., t.Any]) -> None:
self.iter_state["actions"].append(fn)

if self.after is not None:
self.after(retry_state)
def _run_retry(self, retry_state: "RetryCallState") -> None:
self.iter_state["retry_run_result"] = self.retry(retry_state)

def _run_wait(self, retry_state: "RetryCallState") -> None:
if self.wait:
sleep = self.wait(retry_state)
else:
sleep = 0.0

retry_state.upcoming_sleep = sleep

def _run_stop(self, retry_state: "RetryCallState") -> None:
self.statistics["delay_since_first_attempt"] = retry_state.seconds_since_start
if self.stop(retry_state):
self.iter_state["stop_run_result"] = self.stop(retry_state)

def iter(self, retry_state: "RetryCallState") -> t.Union[DoAttempt, DoSleep, t.Any]: # noqa
self._begin_iter(retry_state)
result = None
for action in self.iter_state["actions"]:
result = action(retry_state)
return result

def _begin_iter(self, retry_state: "RetryCallState") -> None: # noqa
self.iter_state.clear()
self.iter_state["actions"] = []

fut = retry_state.outcome
if fut is None:
if self.before is not None:
self._add_action_func(self.before)
self._add_action_func(lambda rs: DoAttempt())
return

self.iter_state["is_explicit_retry"] = fut.failed and isinstance(fut.exception(), TryAgain)
if not self.iter_state["is_explicit_retry"]:
self._add_action_func(self._run_retry)
self._add_action_func(self._post_retry_check_actions)

def _post_retry_check_actions(self, retry_state: "RetryCallState") -> None:
if not (self.iter_state["is_explicit_retry"] or self.iter_state.get("retry_run_result")):
self._add_action_func(lambda rs: rs.outcome.result())
return

if self.after is not None:
self._add_action_func(self.after)

self._add_action_func(self._run_wait)
self._add_action_func(self._run_stop)
self._add_action_func(self._post_stop_check_actions)

def _post_stop_check_actions(self, retry_state: "RetryCallState") -> None:
if self.iter_state["stop_run_result"]:
if self.retry_error_callback:
return self.retry_error_callback(retry_state)
retry_exc = self.retry_error_cls(fut)
if self.reraise:
raise retry_exc.reraise()
raise retry_exc from fut.exception()
self._add_action_func(self.retry_error_callback)
return

def exc_check(rs: "RetryCallState") -> None:
fut = t.cast(Future, rs.outcome)
retry_exc = self.retry_error_cls(fut)
if self.reraise:
raise retry_exc.reraise()
raise retry_exc from fut.exception()

self._add_action_func(exc_check)
return

def next_action(rs: "RetryCallState") -> None:
sleep = rs.upcoming_sleep
rs.next_action = RetryAction(sleep)
rs.idle_for += sleep
self.statistics["idle_for"] += sleep
self.statistics["attempt_number"] += 1

retry_state.next_action = RetryAction(sleep)
retry_state.idle_for += sleep
self.statistics["idle_for"] += sleep
self.statistics["attempt_number"] += 1
self._add_action_func(next_action)

if self.before_sleep is not None:
self.before_sleep(retry_state)
self._add_action_func(self.before_sleep)

return DoSleep(sleep)
self._add_action_func(lambda rs: DoSleep(rs.upcoming_sleep))

def __iter__(self) -> t.Generator[AttemptManager, None, None]:
self.begin()
Expand Down Expand Up @@ -505,16 +557,16 @@ def retry(func: WrappedFn) -> WrappedFn:

@t.overload
def retry(
sleep: t.Callable[[t.Union[int, float]], t.Optional[t.Awaitable[None]]] = sleep,
stop: "StopBaseT" = stop_never,
wait: "WaitBaseT" = wait_none(),
retry: "RetryBaseT" = retry_if_exception_type(),
before: t.Callable[["RetryCallState"], None] = before_nothing,
after: t.Callable[["RetryCallState"], None] = after_nothing,
before_sleep: t.Optional[t.Callable[["RetryCallState"], None]] = None,
sleep: t.Callable[[t.Union[int, float]], t.Union[None, t.Awaitable[None]]] = sleep,
stop: "t.Union[StopBaseT, tasyncio.stop.StopBaseT]" = stop_never,
wait: "t.Union[WaitBaseT, tasyncio.wait.WaitBaseT]" = wait_none(),
retry: "t.Union[RetryBaseT, tasyncio.retry.RetryBaseT]" = retry_if_exception_type(),
before: t.Callable[["RetryCallState"], t.Union[None, t.Awaitable[None]]] = before_nothing,
after: t.Callable[["RetryCallState"], t.Union[None, t.Awaitable[None]]] = after_nothing,
before_sleep: t.Optional[t.Callable[["RetryCallState"], t.Union[None, t.Awaitable[None]]]] = None,
reraise: bool = False,
retry_error_cls: t.Type["RetryError"] = RetryError,
retry_error_callback: t.Optional[t.Callable[["RetryCallState"], t.Any]] = None,
retry_error_callback: t.Optional[t.Callable[["RetryCallState"], t.Union[t.Any, t.Awaitable[t.Any]]]] = None,
) -> t.Callable[[WrappedFn], WrappedFn]:
...

Expand Down Expand Up @@ -549,7 +601,7 @@ def wrap(f: WrappedFn) -> WrappedFn:
return wrap


from tenacity._asyncio import AsyncRetrying # noqa:E402,I100
from tenacity.asyncio import AsyncRetrying # noqa:E402,I100

if tornado:
from tenacity.tornadoweb import TornadoRetrying
Expand Down
94 changes: 0 additions & 94 deletions tenacity/_asyncio.py

This file was deleted.

Loading
Loading