From 6a55573d588e4a817986c8964d97dceb8c46585a Mon Sep 17 00:00:00 2001
From: cyrus durgin <cyrusd@gmail.com>
Date: Thu, 10 Sep 2015 15:57:35 -0700
Subject: [PATCH] Adds two additional parameters to the Retrying() constructor:

    _before_attempts
    _after_attempts

Each of these should be a reference to a function that takes a
single argument, attempt_number.

_before_attempts will be called before each attempt is run;
_after_attempts is run after any *unsuccessful* attempt.

Handy for logging status messages, etc in between attempts.
---
 AUTHORS.rst      |  3 ++-
 retrying.py      | 12 +++++++++++-
 test_retrying.py | 34 ++++++++++++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/AUTHORS.rst b/AUTHORS.rst
index 9cb572f..42a456c 100644
--- a/AUTHORS.rst
+++ b/AUTHORS.rst
@@ -30,4 +30,5 @@ Patches and Suggestions
 - Monty Taylor
 - Maxym Shalenyi
 - Jonathan Herriott
-- Job Evers
\ No newline at end of file
+- Job Evers
+- Cyrus Durgin
diff --git a/retrying.py b/retrying.py
index 3ed312d..79111ae 100644
--- a/retrying.py
+++ b/retrying.py
@@ -68,7 +68,9 @@ class Retrying(object):
                  wrap_exception=False,
                  stop_func=None,
                  wait_func=None,
-                 wait_jitter_max=None):
+                 wait_jitter_max=None,
+                 before_attempts=None,
+                 after_attempts=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
@@ -80,6 +82,8 @@ class Retrying(object):
         self._wait_exponential_multiplier = 1 if wait_exponential_multiplier is None else wait_exponential_multiplier
         self._wait_exponential_max = MAX_WAIT if wait_exponential_max is None else wait_exponential_max
         self._wait_jitter_max = 0 if wait_jitter_max is None else wait_jitter_max
+        self._before_attempts = before_attempts
+        self._after_attempts = after_attempts
 
         # TODO add chaining of stop behaviors
         # stop behavior
@@ -196,6 +200,9 @@ class Retrying(object):
         start_time = int(round(time.time() * 1000))
         attempt_number = 1
         while True:
+            if self._before_attempts:
+                self._before_attempts(attempt_number)
+
             try:
                 attempt = Attempt(fn(*args, **kwargs), attempt_number, False)
             except:
@@ -205,6 +212,9 @@ class Retrying(object):
             if not self.should_reject(attempt):
                 return attempt.get(self._wrap_exception)
 
+            if self._after_attempts:
+                self._after_attempts(attempt_number)
+
             delay_since_first_attempt_ms = int(round(time.time() * 1000)) - start_time
             if self.stop(attempt_number, delay_since_first_attempt_ms):
                 if not self._wrap_exception and attempt.has_exception:
diff --git a/test_retrying.py b/test_retrying.py
index bfef02d..3cc843a 100644
--- a/test_retrying.py
+++ b/test_retrying.py
@@ -434,5 +434,39 @@ class TestDecoratorWrapper(unittest.TestCase):
         self.assertTrue(_retryable_default(NoCustomErrorAfterCount(5)))
         self.assertTrue(_retryable_default_f(NoCustomErrorAfterCount(5)))
 
+class TestBeforeAfterAttempts(unittest.TestCase):
+    _attempt_number = 0
+
+    def test_before_attempts(self):
+        TestBeforeAfterAttempts._attempt_number = 0
+
+        def _before(attempt_number):
+            TestBeforeAfterAttempts._attempt_number = attempt_number
+
+        @retry(wait_fixed = 1000, stop_max_attempt_number = 1, before_attempts = _before)
+        def _test_before():
+            pass
+        
+        _test_before()
+
+        self.assertTrue(TestBeforeAfterAttempts._attempt_number is 1)
+
+    def test_after_attempts(self):
+        TestBeforeAfterAttempts._attempt_number = 0
+
+        def _after(attempt_number):
+            TestBeforeAfterAttempts._attempt_number = attempt_number
+
+        @retry(wait_fixed = 100, stop_max_attempt_number = 3, after_attempts = _after)
+        def _test_after():
+            if TestBeforeAfterAttempts._attempt_number < 2:
+                raise Exception("testing after_attempts handler")
+            else:
+                pass
+        
+        _test_after()
+
+        self.assertTrue(TestBeforeAfterAttempts._attempt_number is 2)
+
 if __name__ == '__main__':
     unittest.main()