diff --git a/retrying.py b/retrying.py index 3ed312d..a2aadc7 100644 --- a/retrying.py +++ b/retrying.py @@ -31,22 +31,24 @@ def retry(*dargs, **dkw): """ # support both @retry and @retry() as valid syntax if len(dargs) == 1 and callable(dargs[0]): + retryer = Retrying() def wrap_simple(f): @six.wraps(f) def wrapped_f(*args, **kw): - return Retrying().call(f, *args, **kw) + return retryer(f, *args, **kw) return wrapped_f return wrap_simple(dargs[0]) else: + retryer = Retrying(*dargs, **dkw) def wrap(f): @six.wraps(f) def wrapped_f(*args, **kw): - return Retrying(*dargs, **dkw).call(f, *args, **kw) + return retryer(f, *args, **kw) return wrapped_f @@ -138,6 +140,17 @@ def __init__(self, self._wrap_exception = wrap_exception + self._attempt_number = 0 + self._start_time = None + + @property + def attempts(self): + return self._attempt_number + + @property + def start_time(self): + return self._start_time + def stop_after_attempt(self, previous_attempt_number, delay_since_first_attempt_ms): """Stop after the previous attempt >= stop_max_attempt_number.""" return previous_attempt_number >= self._stop_max_attempt_number @@ -192,34 +205,34 @@ def should_reject(self, attempt): return reject - def call(self, fn, *args, **kwargs): - start_time = int(round(time.time() * 1000)) - attempt_number = 1 + def __call__(self, fn, *args, **kwargs): + self._start_time = int(round(time.time() * 1000)) + self._attempt_number = 1 while True: try: - attempt = Attempt(fn(*args, **kwargs), attempt_number, False) + attempt = Attempt(fn(*args, **kwargs), self._attempt_number, False) except: tb = sys.exc_info() - attempt = Attempt(tb, attempt_number, True) + attempt = Attempt(tb, self._attempt_number, True) if not self.should_reject(attempt): return attempt.get(self._wrap_exception) - delay_since_first_attempt_ms = int(round(time.time() * 1000)) - start_time - if self.stop(attempt_number, delay_since_first_attempt_ms): + delay_since_first_attempt_ms = int(round(time.time() * 1000)) - self._start_time + if self.stop(self._attempt_number, delay_since_first_attempt_ms): if not self._wrap_exception and attempt.has_exception: # get() on an attempt with an exception should cause it to be raised, but raise just in case raise attempt.get() else: raise RetryError(attempt) else: - sleep = self.wait(attempt_number, delay_since_first_attempt_ms) + sleep = self.wait(self._attempt_number, delay_since_first_attempt_ms) if self._wait_jitter_max: jitter = random.random() * self._wait_jitter_max sleep = sleep + max(0, jitter) time.sleep(sleep / 1000.0) - attempt_number += 1 + self._attempt_number += 1 class Attempt(object): diff --git a/test_retrying.py b/test_retrying.py index bfef02d..ba1cd4c 100644 --- a/test_retrying.py +++ b/test_retrying.py @@ -434,5 +434,33 @@ def test_defaults(self): self.assertTrue(_retryable_default(NoCustomErrorAfterCount(5))) self.assertTrue(_retryable_default_f(NoCustomErrorAfterCount(5))) + +class TestDirectUsage(unittest.TestCase): + def test_direct_call(self): + value = 1 + def working(): + return value + + r = Retrying() + self.assertFalse(r.start_time) + self.assertFalse(r.attempts) + result = r(working) + self.assertEqual(value, result) + self.assertEqual(1, r.attempts) + + def test_direct_call_with_args(self): + value = 1 + def working(value): + return value + + r = Retrying() + self.assertFalse(r.start_time) + self.assertFalse(r.attempts) + result = r(working, value) + self.assertEqual(value, result) + self.assertEqual(1, r.attempts) + self.assertTrue(r.start_time) + + if __name__ == '__main__': unittest.main()