diff options
Diffstat (limited to 'retrying.py')
-rw-r--r-- | retrying.py | 14 |
1 files changed, 11 insertions, 3 deletions
diff --git a/retrying.py b/retrying.py index ccf9668..50fc439 100644 --- a/retrying.py +++ b/retrying.py @@ -108,7 +108,9 @@ class Retrying(object): wait_exponential_multiplier=None, wait_exponential_max=None, retry_on_exception=None, retry_on_result=None, - wrap_exception=False): + wrap_exception=False, + stop_func=None, + wait_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 @@ -129,7 +131,10 @@ class Retrying(object): if stop_max_delay is not None: stop_funcs.append(self.stop_after_delay) - if stop is None: + if stop_func is not None: + self.stop = stop_func + + elif stop is None: self.stop = lambda attempts, delay: any(f(attempts, delay) for f in stop_funcs) else: @@ -150,7 +155,10 @@ class Retrying(object): if wait_exponential_multiplier is not None or wait_exponential_max is not None: wait_funcs.append(self.exponential_sleep) - if wait is None: + if wait_func is not None: + self.wait = wait_func + + elif wait is None: self.wait = lambda attempts, delay: max(f(attempts, delay) for f in wait_funcs) else: |