diff options
| -rw-r--r-- | .gitignore | 3 | ||||
| -rw-r--r-- | AUTHORS.rst | 3 | ||||
| -rw-r--r-- | retrying.py | 25 | ||||
| -rw-r--r-- | test_retrying.py | 40 |
4 files changed, 64 insertions, 7 deletions
@@ -1,4 +1,5 @@ .idea dist *.pyc -*.egg-info
\ No newline at end of file +*.egg-info +build 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 547509f..bcb7a9d 100644 --- a/retrying.py +++ b/retrying.py @@ -23,6 +23,12 @@ import traceback MAX_WAIT = 1073741823 +def _retry_if_exception_of_type(retryable_types): + def _retry_if_exception_these_types(exception): + return isinstance(exception, retryable_types) + return _retry_if_exception_these_types + + def retry(*dargs, **dkw): """ Decorator function that instantiates the Retrying object @@ -69,7 +75,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 @@ -82,6 +90,8 @@ class Retrying(object): self._wait_exponential_max = MAX_WAIT if wait_exponential_max is None else wait_exponential_max self._wait_incrementing_max = MAX_WAIT if wait_incrementing_max is None else wait_incrementing_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 @@ -129,9 +139,14 @@ class Retrying(object): if retry_on_exception is None: self._retry_on_exception = self.always_reject else: + # this allows for providing a tuple of exception types that + # should be allowed to retry on, and avoids having to create + # a callback that does the same thing + if isinstance(retry_on_exception, (tuple)): + retry_on_exception = _retry_if_exception_of_type( + retry_on_exception) self._retry_on_exception = retry_on_exception - # TODO simplify retrying by Exception types # retry on result filter if retry_on_result is None: self._retry_on_result = self.never_reject @@ -203,6 +218,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: @@ -212,6 +230,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..8ce4ac3 100644 --- a/test_retrying.py +++ b/test_retrying.py @@ -252,7 +252,7 @@ def _retryable_test_with_stop(thing): return thing.go() -@retry(retry_on_exception=retry_if_exception_of_type(IOError)) +@retry(retry_on_exception=(IOError,)) def _retryable_test_with_exception_type_io(thing): return thing.go() @@ -264,14 +264,14 @@ def _retryable_test_with_exception_type_io_wrap(thing): @retry( stop_max_attempt_number=3, - retry_on_exception=retry_if_exception_of_type(IOError)) + retry_on_exception=(IOError,)) def _retryable_test_with_exception_type_io_attempt_limit(thing): return thing.go() @retry( stop_max_attempt_number=3, - retry_on_exception=retry_if_exception_of_type(IOError), + retry_on_exception=(IOError,), wrap_exception=True) def _retryable_test_with_exception_type_io_attempt_limit_wrap(thing): return thing.go() @@ -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() |