Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# Copyright 2016 Étienne Bersac
2# Copyright 2016 Julien Danjou
3# Copyright 2016 Joshua Harlow
4# Copyright 2013-2014 Ray Holder
5#
6# Licensed under the Apache License, Version 2.0 (the "License");
7# you may not use this file except in compliance with the License.
8# You may obtain a copy of the License at
9#
10# http://www.apache.org/licenses/LICENSE-2.0
11#
12# Unless required by applicable law or agreed to in writing, software
13# distributed under the License is distributed on an "AS IS" BASIS,
14# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15# See the License for the specific language governing permissions and
16# limitations under the License.
18import functools
19import sys
20import typing as t
22import tenacity
23from tenacity import (
24 AttemptManager,
25 BaseRetrying,
26 DoAttempt,
27 DoSleep,
28 RetryCallState,
29 RetryError,
30 _RetryDecorated,
31 _utils,
32 after_nothing,
33 before_nothing,
34)
35from tenacity._utils import override
37# Import all built-in retry strategies for easier usage.
38from .retry import (
39 RetryBaseT,
40 retry_all,
41 retry_any,
42 retry_if_exception,
43 retry_if_result,
44)
46if t.TYPE_CHECKING:
47 from tenacity.retry import RetryBaseT as SyncRetryBaseT
48 from tenacity.stop import StopBaseT
49 from tenacity.wait import WaitBaseT
51WrappedFnReturnT = t.TypeVar("WrappedFnReturnT")
52WrappedFn = t.TypeVar("WrappedFn", bound=t.Callable[..., t.Awaitable[t.Any]])
53P = t.ParamSpec("P")
54R = t.TypeVar("R")
57def _portable_async_sleep(seconds: float) -> t.Awaitable[None]:
58 # If trio is already imported, then importing it is cheap.
59 # If trio isn't already imported, then it's definitely not running, so we
60 # can skip further checks.
61 if "trio" in sys.modules:
62 # If trio is available, then sniffio is too
63 import sniffio
64 import trio
66 if sniffio.current_async_library() == "trio":
67 return trio.sleep(seconds) # noqa: ASYNC105
68 # Otherwise, assume asyncio
69 # Lazy import asyncio as it's expensive (responsible for 25-50% of total import overhead).
70 import asyncio
72 return asyncio.sleep(seconds)
75class AsyncRetrying(BaseRetrying):
76 def __init__(
77 self,
78 sleep: t.Callable[
79 [int | float], t.Awaitable[None] | None
80 ] = _portable_async_sleep,
81 stop: "StopBaseT" = tenacity.stop.stop_never,
82 wait: "WaitBaseT" = tenacity.wait.wait_none(),
83 retry: "SyncRetryBaseT | RetryBaseT" = tenacity.retry_if_exception_type(),
84 before: t.Callable[
85 ["RetryCallState"], t.Awaitable[None] | None
86 ] = before_nothing,
87 after: t.Callable[["RetryCallState"], t.Awaitable[None] | None] = after_nothing,
88 before_sleep: t.Callable[["RetryCallState"], t.Awaitable[None] | None]
89 | None = None,
90 reraise: bool = False,
91 retry_error_cls: type["RetryError"] = RetryError,
92 retry_error_callback: t.Callable[["RetryCallState"], t.Any | t.Awaitable[t.Any]]
93 | None = None,
94 name: str | None = None,
95 enabled: bool = True,
96 ) -> None:
97 super().__init__(
98 sleep=sleep, # type: ignore[arg-type]
99 stop=stop,
100 wait=wait,
101 retry=retry, # type: ignore[arg-type]
102 before=before, # type: ignore[arg-type]
103 after=after, # type: ignore[arg-type]
104 before_sleep=before_sleep, # type: ignore[arg-type]
105 reraise=reraise,
106 retry_error_cls=retry_error_cls,
107 retry_error_callback=retry_error_callback,
108 name=name,
109 enabled=enabled,
110 )
112 @override
113 async def __call__( # type: ignore[override]
114 self, fn: WrappedFn, *args: t.Any, **kwargs: t.Any
115 ) -> WrappedFnReturnT:
116 is_async = _utils.is_coroutine_callable(fn)
117 if not self.enabled:
118 if is_async:
119 return await fn(*args, **kwargs) # type: ignore[no-any-return]
120 return fn(*args, **kwargs) # type: ignore[return-value]
122 self.begin()
124 retry_state = RetryCallState(retry_object=self, fn=fn, args=args, kwargs=kwargs)
125 while True:
126 do = await self.iter(retry_state=retry_state)
127 if isinstance(do, DoAttempt):
128 try:
129 if is_async:
130 result = await fn(*args, **kwargs)
131 else:
132 result = fn(*args, **kwargs)
133 except BaseException:
134 retry_state.set_exception(sys.exc_info()) # type: ignore[arg-type]
135 else:
136 retry_state.set_result(result)
137 elif isinstance(do, DoSleep):
138 retry_state.prepare_for_next_attempt()
139 await self.sleep(do) # type: ignore[misc]
140 else:
141 return do # type: ignore[no-any-return]
143 @override
144 def _add_action_func(self, fn: t.Callable[..., t.Any]) -> None:
145 self.iter_state.actions.append(_utils.wrap_to_async_func(fn))
147 @override
148 async def _run_retry(self, retry_state: "RetryCallState") -> None: # type: ignore[override]
149 self.iter_state.retry_run_result = await _utils.wrap_to_async_func(self.retry)(
150 retry_state
151 )
153 @override
154 async def _run_wait(self, retry_state: "RetryCallState") -> None: # type: ignore[override]
155 retry_state.upcoming_sleep = await _utils.wrap_to_async_func(self.wait)(
156 retry_state
157 )
159 @override
160 async def _run_stop(self, retry_state: "RetryCallState") -> None: # type: ignore[override]
161 self.statistics["delay_since_first_attempt"] = retry_state.seconds_since_start
162 self.iter_state.stop_run_result = await _utils.wrap_to_async_func(self.stop)(
163 retry_state
164 )
166 @override
167 async def iter(self, retry_state: "RetryCallState") -> DoAttempt | DoSleep | t.Any:
168 self._begin_iter(retry_state)
169 result = None
170 for action in self.iter_state.actions:
171 result = await action(retry_state)
172 return result
174 @override
175 def __iter__(self) -> t.Generator[AttemptManager, None, None]:
176 raise TypeError("AsyncRetrying object is not iterable")
178 def __aiter__(self) -> "AsyncRetrying":
179 if not self.enabled:
180 self._retry_state = RetryCallState(self, fn=None, args=(), kwargs={})
181 self._disabled_iter_done = False
182 return self
184 self.begin()
185 self._retry_state = RetryCallState(self, fn=None, args=(), kwargs={})
186 return self
188 async def __anext__(self) -> AttemptManager:
189 if not self.enabled:
190 if self._disabled_iter_done:
191 outcome = self._retry_state.outcome
192 if outcome is not None and outcome.failed:
193 raise outcome.exception() # type: ignore[misc]
194 raise StopAsyncIteration
195 self._disabled_iter_done = True
196 return AttemptManager(retry_state=self._retry_state)
198 while True:
199 do = await self.iter(retry_state=self._retry_state)
200 if do is None:
201 raise StopAsyncIteration
202 if isinstance(do, DoAttempt):
203 return AttemptManager(retry_state=self._retry_state)
204 if isinstance(do, DoSleep):
205 self._retry_state.prepare_for_next_attempt()
206 await self.sleep(do) # type: ignore[misc]
207 else:
208 raise StopAsyncIteration
210 @override
211 def wraps(self, fn: t.Callable[P, R]) -> _RetryDecorated[P, R]:
212 wrapped = super().wraps(fn)
213 # Ensure wrapper is recognized as a coroutine function.
215 @functools.wraps(
216 fn, functools.WRAPPER_ASSIGNMENTS + ("__defaults__", "__kwdefaults__")
217 )
218 async def async_wrapped(*args: t.Any, **kwargs: t.Any) -> t.Any:
219 if not self.enabled:
220 return await fn(*args, **kwargs) # type: ignore[misc]
221 # Always create a copy to prevent overwriting the local contexts when
222 # calling the same wrapped functions multiple times in the same stack
223 copy = self.copy()
224 # Reuse the same statistics dict rather than rebinding the attribute
225 # so that the stats stay visible through additional decorators that
226 # copy attributes via functools.wraps (which copies the reference to
227 # this dict into the outer wrapper's __dict__). See issue #519.
228 stats = async_wrapped.statistics # type: ignore[attr-defined]
229 stats.clear()
230 copy._local.statistics = stats # noqa: SLF001
231 self._local.statistics = stats
232 return await copy(fn, *args, **kwargs) # type: ignore[type-var]
234 # Preserve attributes
235 async_wrapped.retry = self # type: ignore[attr-defined]
236 async_wrapped.retry_with = wrapped.retry_with # type: ignore[attr-defined]
237 async_wrapped.statistics = {} # type: ignore[attr-defined]
239 return t.cast("_RetryDecorated[P, R]", async_wrapped)
242__all__ = [
243 "AsyncRetrying",
244 "WrappedFn",
245 "retry_all",
246 "retry_any",
247 "retry_if_exception",
248 "retry_if_result",
249]