Allow chaining of wait lengths through customized aggregation function

AUTHORS.rst

@@ -32,3 +32,4 @@ Patches and Suggestions
 - Jonathan Herriott
 - Job Evers
 - Cyrus Durgin
+- Justin Patrin

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)

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")