diff --git a/AUTHORS.rst b/AUTHORS.rst index 42a456c..c2b6b18 100644 --- a/AUTHORS.rst +++ b/AUTHORS.rst @@ -32,3 +32,4 @@ Patches and Suggestions - Jonathan Herriott - Job Evers - Cyrus Durgin +- Justin Patrin diff --git a/retrying.py b/retrying.py index bcb7a9d..f83c9a7 100644 --- a/retrying.py +++ b/retrying.py @@ -77,7 +77,8 @@ def __init__(self, wait_func=None, wait_jitter_max=None, before_attempts=None, - after_attempts=None): + after_attempts=None, + wait_aggregation_func=None): self._stop_max_attempt_number = 5 if stop_max_attempt_number is None else stop_max_attempt_number self._stop_max_delay = 100 if stop_max_delay is None else stop_max_delay @@ -111,7 +112,6 @@ def __init__(self, else: self.stop = getattr(self, stop) - # TODO add chaining of wait behaviors # wait behavior wait_funcs = [lambda *args, **kwargs: 0] if wait_fixed is not None: @@ -130,7 +130,8 @@ def __init__(self, self.wait = wait_func elif wait is None: - self.wait = lambda attempts, delay: max(f(attempts, delay) for f in wait_funcs) + wait_aggregation_func = max if wait_aggregation_func is None else wait_aggregation_func + self.wait = lambda attempts, delay: wait_aggregation_func(f(attempts, delay) for f in wait_funcs) else: self.wait = getattr(self, wait) diff --git a/test_retrying.py b/test_retrying.py index 8ce4ac3..5e4be47 100644 --- a/test_retrying.py +++ b/test_retrying.py @@ -123,6 +123,18 @@ def test_exponential_with_max_wait_and_multiplier(self): self.assertEqual(r.wait(7, 0), 50000) self.assertEqual(r.wait(50, 0), 50000) + def test_wait_aggregation_func(self): + r = Retrying(wait_exponential_max=50000, wait_exponential_multiplier=1000, wait_fixed=1, + wait_aggregation_func=sum) + self.assertEqual(r.wait(1, 0), 2001) + self.assertEqual(r.wait(2, 0), 4001) + self.assertEqual(r.wait(3, 0), 8001) + self.assertEqual(r.wait(4, 0), 16001) + self.assertEqual(r.wait(5, 0), 32001) + self.assertEqual(r.wait(6, 0), 50001) + self.assertEqual(r.wait(7, 0), 50001) + self.assertEqual(r.wait(50, 0), 50001) + def test_legacy_explicit_wait_type(self): Retrying(wait="exponential_sleep")