feat(wait): add wait_exception strategy

Add `wait_exception` which calls a user-defined predicate with the caught
exception to determine the wait time.

Raises `RuntimeError` if the outcome has no exception.

Author: capitan-davide

tenacity/__init__.py

@@ -59,6 +59,7 @@
 # Import all built-in wait strategies for easier usage.
 from .wait import wait_chain  # noqa
 from .wait import wait_combine  # noqa
+from .wait import wait_exception  # noqa
 from .wait import wait_exponential  # noqa
 from .wait import wait_fixed  # noqa
 from .wait import wait_incrementing  # noqa
@@ -686,6 +687,7 @@ def wrap(f: WrappedFn) -> WrappedFn:
     "stop_when_event_set",
     "wait_chain",
     "wait_combine",
+    "wait_exception",
     "wait_exponential",
     "wait_fixed",
     "wait_incrementing",

tenacity/wait.py

@@ -113,6 +113,42 @@ def __call__(self, retry_state: "RetryCallState") -> float:
         return wait_func(retry_state=retry_state)
 
 
+class wait_exception(wait_base):
+    """Wait strategy that waits the amount of time returned by the predicate.
+
+    The predicate is passed the exception object. Based on the exception, the
+    user can decide how much time to wait before retrying.
+
+    For example::
+
+        def http_error(exception: BaseException) -> float:
+            if (
+                isinstance(exception, requests.HTTPError)
+                and exception.response.status_code == requests.codes.too_many_requests
+            ):
+                return float(exception.response.headers.get("Retry-After", "1"))
+            return 60.0
+
+
+        @retry(
+            stop=stop_after_attempt(3),
+            wait=wait_exception(http_error),
+        )
+        def http_get_request(url: str) -> None:
+            response = requests.get(url)
+            response.raise_for_status()
+    """
+
+    def __init__(self, predicate: typing.Callable[[BaseException], float]) -> None:
+        self.predicate = predicate
+
+    def __call__(self, retry_state: "RetryCallState") -> float:
+        exception = retry_state.outcome.exception()
+        if exception is None:
+            raise RuntimeError("outcome failed but the exception is None")
+        return self.predicate(exception)
+
+
 class wait_incrementing(wait_base):
     """Wait an incremental amount of time after each attempt.
 

tests/test_tenacity.py

@@ -369,6 +369,24 @@ def test_wait_combine(self):
             self.assertLess(w, 8)
             self.assertGreaterEqual(w, 5)
 
+    def test_wait_exception(self):
+        def predicate(exc):
+            if isinstance(exc, ValueError):
+                return 3.5
+            return 10.0
+
+        r = Retrying(wait=tenacity.wait_exception(predicate))
+
+        fut1 = tenacity.Future.construct(1, ValueError(), True)
+        self.assertEqual(r.wait(make_retry_state(1, 0, last_result=fut1)), 3.5)
+
+        fut2 = tenacity.Future.construct(1, KeyError(), True)
+        self.assertEqual(r.wait(make_retry_state(1, 0, last_result=fut2)), 10.0)
+
+        fut3 = tenacity.Future.construct(1, None, False)
+        with self.assertRaises(RuntimeError):
+            r.wait(make_retry_state(1, 0, last_result=fut3))
+
     def test_wait_double_sum(self):
         r = Retrying(wait=tenacity.wait_random(0, 3) + tenacity.wait_fixed(5))
         # Test it a few time since it's random