Coverage for src / toolbox_python / retry.py: 100%
75 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-02 22:56 +0000
1# ============================================================================ #
2# #
3# Title : Retry #
4# Purpose : Automatically retry a given function when a specific #
5# `Exception` is thrown. #
6# #
7# ============================================================================ #
10# ---------------------------------------------------------------------------- #
11# #
12# Overview ####
13# #
14# ---------------------------------------------------------------------------- #
17# ---------------------------------------------------------------------------- #
18# Description ####
19# ---------------------------------------------------------------------------- #
22"""
23!!! note "Summary"
24 The `retry` module is for enabling automatic retrying of a given function when a specific `Exception` is thrown.
25"""
28# ---------------------------------------------------------------------------- #
29# #
30# Setup ####
31# #
32# ---------------------------------------------------------------------------- #
35# ---------------------------------------------------------------------------- #
36# Imports ####
37# ---------------------------------------------------------------------------- #
40# ## Python StdLib Imports ----
41import inspect
42import logging
43from builtins import Exception
44from functools import wraps
45from logging import Logger
46from time import sleep
47from types import ModuleType
48from typing import Any, Callable, Literal, NoReturn, Optional, TypeVar, Union, overload
50# ## Python Third Party Imports ----
51from typeguard import typechecked
53# ## Local First Party Imports ----
54from toolbox_python.checkers import assert_is_valid
55from toolbox_python.classes import get_full_class_name
56from toolbox_python.output import print_or_log_output
59# ---------------------------------------------------------------------------- #
60# Exports ####
61# ---------------------------------------------------------------------------- #
64__all__: list[str] = ["retry"]
67# ---------------------------------------------------------------------------- #
68# Types ####
69# ---------------------------------------------------------------------------- #
72_exceptions = Union[
73 type[Exception],
74 list[type[Exception]],
75 tuple[type[Exception], ...],
76]
77"""
78!!! note "Summary"
79 A type alias for a single or collection of `Exception` types.
80"""
83R = TypeVar("R")
84"""
85!!! note "Summary"
86 A generic type variable to represent the return type of a function.
87"""
90# ---------------------------------------------------------------------------- #
91# #
92# Classes ####
93# #
94# ---------------------------------------------------------------------------- #
97class _Retry:
98 """
99 !!! note "Summary"
100 A helper class to handle the retry logic for the `retry` decorator.
102 ???+ abstract "Details"
103 This class is not intended to be used directly. Instead, it is used internally by the `retry` decorator to manage the retry logic.
105 Methods:
106 run(): Run the retry loop for the given function.
107 """
109 def __init__(
110 self,
111 exceptions: _exceptions,
112 tries: int,
113 delay: int,
114 print_or_log: Literal["print", "log"],
115 log: Optional[Logger],
116 ) -> None:
117 """
118 !!! note "Summary"
119 Initialize the `_Retry` class with the given parameters.
121 Params:
122 exceptions (_exceptions):
123 A given single or collection of expected exceptions for which to catch and retry for.
124 tries (int):
125 The number of retries to attempt.
126 delay (int):
127 The number of seconds to delay between each retry.
128 print_or_log (Literal["print", "log"]):
129 Whether or not the messages should be written to the terminal in a `#!py print()` statement, or to a log file in a `#!py log()` statement.
130 log (Optional[Logger]):
131 An optional logger instance to use when `print_or_log` is set to `"log"`.
132 """
133 self.exceptions: tuple[type[Exception], ...] = (
134 tuple(exceptions) if isinstance(exceptions, (list, tuple)) else (exceptions,)
135 )
136 self.tries: int = tries
137 self.delay: int = delay
138 self.print_or_log: Literal["print", "log"] = print_or_log
139 self.log: Optional[Logger] = log
141 def run(self, func: Callable[..., R], *args: Any, **kwargs: Any) -> R:
142 """
143 !!! note "Summary"
144 Run the retry loop for the given function.
145 """
146 for i in range(1, self.tries + 1):
147 try:
148 results = func(*args, **kwargs)
149 self._handle_success(i)
150 return results
151 except self.exceptions as e:
152 self._handle_expected_error(i, e)
153 except Exception as exc:
154 self._handle_unexpected_error(i, exc)
155 self._handle_final_failure()
157 def _handle_success(self, i: int) -> None:
158 message: str = f"Successfully executed at iteration {i}."
159 print_or_log_output(
160 message=message,
161 print_or_log=self.print_or_log,
162 log=self.log,
163 log_level="info",
164 )
166 def _handle_expected_error(self, i: int, e: Exception) -> None:
167 message = (
168 f"Caught an expected error at iteration {i}: "
169 f"`{get_full_class_name(e)}`. "
170 f"Retrying in {self.delay} seconds..."
171 )
172 print_or_log_output(
173 message=message,
174 print_or_log=self.print_or_log,
175 log=self.log,
176 log_level="warning",
177 )
178 sleep(self.delay)
180 def _handle_unexpected_error(self, i: int, exc: Exception) -> None:
181 excs = self.exceptions if isinstance(self.exceptions, (list, tuple)) else (self.exceptions,)
182 exc_names: list[str] = [e.__name__ for e in excs]
183 if any(name in f"{exc}" for name in exc_names):
184 caught_errors: list[str] = [name for name in exc_names if name in f"{exc}"]
185 message: str = (
186 f"Caught an unexpected, known error at iteration {i}: "
187 f"`{get_full_class_name(exc)}`.\n"
188 f"Who's message contains reference to underlying exception(s): {caught_errors}.\n"
189 f"Retrying in {self.delay} seconds..."
190 )
191 print_or_log_output(
192 message=message,
193 print_or_log=self.print_or_log,
194 log=self.log,
195 log_level="warning",
196 )
197 sleep(self.delay)
198 else:
199 message = f"Caught an unexpected error at iteration {i}: `{get_full_class_name(exc)}`."
200 print_or_log_output(
201 message=message,
202 print_or_log=self.print_or_log,
203 log=self.log,
204 log_level="error",
205 )
206 raise RuntimeError(message) from exc
208 def _handle_final_failure(self) -> NoReturn:
209 message: str = f"Still could not write after {self.tries} iterations. Please check."
210 print_or_log_output(
211 message=message,
212 print_or_log=self.print_or_log,
213 log=self.log,
214 log_level="error",
215 )
216 raise RuntimeError(message)
219# ---------------------------------------------------------------------------- #
220# #
221# Functions ####
222# #
223# ---------------------------------------------------------------------------- #
226@overload
227@typechecked
228def retry(
229 exceptions: _exceptions = Exception,
230 tries: int = 1,
231 delay: int = 0,
232 print_or_log: Literal["print"] = "print",
233) -> Callable[[Callable[..., R]], Callable[..., R]]: ...
234@overload
235@typechecked
236def retry(
237 exceptions: _exceptions = Exception,
238 tries: int = 1,
239 delay: int = 0,
240 print_or_log: Literal["log"] = "log",
241) -> Callable[[Callable[..., R]], Callable[..., R]]: ...
242@typechecked
243def retry(
244 exceptions: _exceptions = Exception,
245 tries: int = 1,
246 delay: int = 0,
247 print_or_log: Literal["print", "log"] = "print",
248) -> Callable[[Callable[..., R]], Callable[..., R]]:
249 """
250 !!! note "Summary"
251 Retry a given function a number of times. Catching any known exceptions when they are given. And returning any output to either a terminal or a log file.
253 !!! deprecation "Deprecated"
254 This function is deprecated. Please use the [`retry()`][func] decorator from the [`stamina`][docs] package instead.<br>
255 For more info, see: [Docs][docs], [GitHub][github], [PyPi][pypi].
256 [func]: https://stamina.hynek.me/en/stable/api.html#stamina.retry
257 [docs]: https://stamina.hynek.me/en/stable/index.html
258 [github]: https://github.com/hynek/stamina/
259 [pypi]: https://pypi.org/project/stamina/
261 ???+ abstract "Details"
262 This function should always be implemented as a decorator.<br>
263 It is written based on the premise that a certain process may fail and return a given message, but that is known and expected, and you just want to wait a second or so then retry again.<br>
264 Typically, this is seen in async processes, or when writing data to a `delta` table when there may be concurrent read/writes occurring at the same time. In these instances, you will know the error message and can re-try again a certain number of times.
266 Params:
267 exceptions (_exceptions, optional):
268 A given single or collection of expected exceptions for which to catch and retry for.<br>
269 Defaults to `#!py Exception`.
270 tries (int, optional):
271 The number of retries to attempt. If the underlying process is still failing after this number of attempts, then throw a hard error and alert the user.<br>
272 Defaults to `#!py 1`.
273 delay (int, optional):
274 The number of seconds to delay between each retry.<br>
275 Defaults to `#!py 0`.
276 print_or_log (Literal["print", "log"], optional):
277 Whether or not the messages should be written to the terminal in a `#!py print()` statement, or to a log file in a `#!py log()` statement.<br>
278 Defaults to `#!py "print"`.
280 Raises:
281 (TypeCheckError):
282 If any of the inputs parsed to the parameters of this function are not the correct type. Uses the [`@typeguard.typechecked`](https://typeguard.readthedocs.io/en/stable/api.html#typeguard.typechecked) decorator.
283 (ValueError):
284 If either `tries` or `delay` are less than `#!py 0`
285 (RuntimeError):
286 If _either_ an unexpected `#!py Exception` was thrown, which was not declared in the `exceptions` collection, _or_ if the `func` was still not able to be executed after `tries` number of iterations.
288 Returns:
289 result (Optional[Any]):
290 The result from the underlying function, if any.
292 ???+ example "Examples"
294 ```pycon {.py .python linenums="1" title="Imports"}
295 >>> from toolbox_python.retry import retry
296 ```
298 ```{.py .python linenums="1" title="Example 1: No error"}
299 >>> @retry(tries=5, delay=1, print_or_log="print")
300 >>> def simple_func(var1: str = "this") -> str:
301 ... return var1
302 ...
303 >>> simple_func()
304 ```
305 <div class="result" markdown>
306 ```{.sh .shell title="Terminal"}
307 # No error
308 ```
309 </div>
311 ```{.py .python linenums="1" title="Example 2: Expected error"}
312 >>> @retry(exceptions=TypeError, tries=5, delay=1, print_or_log="print")
313 >>> def failing_func(var1: str = "that") -> None:
314 ... raise ValueError("Incorrect value")
315 ...
316 >>> failing_func()
317 ```
318 <div class="result" markdown>
319 ```{.sh .shell title="Terminal"}
320 Caught an expected error at iteration 1: `ValueError`. Retrying in 1 seconds...
321 Caught an expected error at iteration 2: `ValueError`. Retrying in 1 seconds...
322 Caught an expected error at iteration 3: `ValueError`. Retrying in 1 seconds...
323 Caught an expected error at iteration 4: `ValueError`. Retrying in 1 seconds...
324 Caught an expected error at iteration 5: `ValueError`. Retrying in 1 seconds...
325 RuntimeError: Still could not write after 5 iterations. Please check.
326 ```
327 </div>
329 ??? success "Credit"
330 Inspiration from:
332 - https://pypi.org/project/retry/
333 - https://stackoverflow.com/questions/21786382/pythonic-way-of-retry-running-a-function#answer-21788594
334 """
336 assert_is_valid(tries, ">=", 0)
337 assert_is_valid(delay, ">=", 0)
339 exceptions = tuple(exceptions) if isinstance(exceptions, (list, tuple)) else (exceptions,)
341 log: Optional[Logger] = None
343 if print_or_log == "log":
344 stk: inspect.FrameInfo = inspect.stack()[2]
345 mod: Union[ModuleType, None] = inspect.getmodule(stk[0])
346 if mod is not None:
347 log: Optional[Logger] = logging.getLogger(mod.__name__)
349 def decorator(func: Callable[..., R]) -> Callable[..., R]:
350 @wraps(func)
351 def wrapper(*args: Any, **kwargs: Any) -> R:
352 retry_handler = _Retry(
353 exceptions=exceptions,
354 tries=tries,
355 delay=delay,
356 print_or_log=print_or_log,
357 log=log,
358 )
359 return retry_handler.run(func, *args, **kwargs)
361 return wrapper
363 return decorator