diff --git a/README.rst b/README.rst index e1d0e9c..9ee0dfd 100644 --- a/README.rst +++ b/README.rst @@ -38,7 +38,7 @@ retry decorator .. code:: python - def retry(exceptions=Exception, tries=-1, delay=0, max_delay=None, backoff=1, jitter=0, logger=logging_logger): + def retry(exceptions=Exception, tries=-1, delay=0, max_delay=None, backoff=1, jitter=0, logger=logging_logger, generator=False): """Return a retry decorator. :param exceptions: an exception or a tuple of exceptions to catch. default: Exception. @@ -50,6 +50,8 @@ retry decorator fixed if a number, random if a range tuple (min, max) :param logger: logger.warning(fmt, error, delay) will be called on failed attempts. default: retry.logging_logger. if None, logging is disabled. + :param generator: if True, assumes decorated function returns a generator and wraps + it in a new generator. Will retry from start on failure. """ Various retrying logic can be achieved by combination of arguments. @@ -101,6 +103,29 @@ Examples logging.basicConfig() make_trouble() +.. code:: python + vals = [RuntimeError(0), 1, 2, RuntimeError(3), 4] + + @retry(generator=True) + def make_trouble(): + for v in vals: + if isinstance(v, BaseException): + vals.remove(v) + raise v + else: + yield v + + if __name__ == '__main__': + # [1, 2, 1, 2, 4] + # Actually: + # + # + # 1, 2 + # + # 1, 2, 4 + print([x for x in make_trouble()]) + + retry_call ^^^^^^^^^^ @@ -158,3 +183,29 @@ This is very similar to the decorator, except that it takes a function and its a +retry_generator +^^^^^^^^^^ + +.. code:: python + + def retry_generator(f, fargs=None, fkwargs=None, exceptions=Exception, tries=-1, delay=0, max_delay=None, backoff=1, jitter=0, + logger=logging_logger): + """ + Return a generator that wraps a generator returned by a given function, retrying from the start on failure. + + :param f: the function that returns a generator + :param fargs: the positional arguments of the function to execute. + :param fkwargs: the named arguments of the function to execute. + :param exceptions: an exception or a tuple of exceptions to catch. default: Exception. + :param tries: the maximum number of attempts. default: -1 (infinite). + :param delay: initial delay between attempts. default: 0. + :param max_delay: the maximum value of delay. default: None (no limit). + :param backoff: multiplier applied to delay between attempts. default: 1 (no backoff). + :param jitter: extra seconds added to delay between attempts. default: 0. + fixed if a number, random if a range tuple (min, max) + :param logger: logger.warning(fmt, error, delay) will be called on failed attempts. + default: retry.logging_logger. if None, logging is disabled. + :returns: the result of the f function. + """ + +This exists for the same reason as retry_call: to be able to dynamically adjust the retry arguments. diff --git a/retry/api.py b/retry/api.py index 4a404b9..8354041 100644 --- a/retry/api.py +++ b/retry/api.py @@ -1,6 +1,7 @@ import logging import random import time +import types from functools import partial @@ -51,7 +52,52 @@ def __retry_internal(f, exceptions=Exception, tries=-1, delay=0, max_delay=None, _delay = min(_delay, max_delay) -def retry(exceptions=Exception, tries=-1, delay=0, max_delay=None, backoff=1, jitter=0, logger=logging_logger): +def __retry_internal_generator(f, exceptions=Exception, tries=-1, delay=0, max_delay=None, backoff=1, jitter=0, + logger=logging_logger): + """ + Iterate through a generator returned by a given function and retry from the start if it failed. + + :param f: the function that returns a generator + :param exceptions: an exception or a tuple of exceptions to catch. default: Exception. + :param tries: the maximum number of attempts. default: -1 (infinite). + :param delay: initial delay between attempts. default: 0. + :param max_delay: the maximum value of delay. default: None (no limit). + :param backoff: multiplier applied to delay between attempts. default: 1 (no backoff). + :param jitter: extra seconds added to delay between attempts. default: 0. + fixed if a number, random if a range tuple (min, max) + :param logger: logger.warning(fmt, error, delay) will be called on failed attempts. + default: retry.logging_logger. if None, logging is disabled. + :returns: the result of the f function. + """ + _tries, _delay = tries, delay + while _tries: + try: + for x in f(): + yield x + raise StopIteration + except StopIteration: + raise + except exceptions as e: + _tries -= 1 + if not _tries: + raise + + if logger is not None: + logger.warning('%s, retrying in %s seconds...', e, _delay) + + time.sleep(_delay) + _delay *= backoff + + if isinstance(jitter, tuple): + _delay += random.uniform(*jitter) + else: + _delay += jitter + + if max_delay is not None: + _delay = min(_delay, max_delay) + + +def retry(exceptions=Exception, tries=-1, delay=0, max_delay=None, backoff=1, jitter=0, logger=logging_logger, generator=False): """Returns a retry decorator. :param exceptions: an exception or a tuple of exceptions to catch. default: Exception. @@ -63,6 +109,8 @@ def retry(exceptions=Exception, tries=-1, delay=0, max_delay=None, backoff=1, ji fixed if a number, random if a range tuple (min, max) :param logger: logger.warning(fmt, error, delay) will be called on failed attempts. default: retry.logging_logger. if None, logging is disabled. + :param generator: if True, assumes decorated function returns a generator and wraps + it in a new generator. Will retry from start on failure. :returns: a retry decorator. """ @@ -70,8 +118,13 @@ def retry(exceptions=Exception, tries=-1, delay=0, max_delay=None, backoff=1, ji def retry_decorator(f, *fargs, **fkwargs): args = fargs if fargs else list() kwargs = fkwargs if fkwargs else dict() - return __retry_internal(partial(f, *args, **kwargs), exceptions, tries, delay, max_delay, backoff, jitter, - logger) + pf = partial(f, *args, **kwargs) + if generator: + return __retry_internal_generator(pf, exceptions, tries, delay, max_delay, backoff, jitter, + logger) + else: + return __retry_internal(pf, exceptions, tries, delay, max_delay, backoff, jitter, + logger) return retry_decorator @@ -99,3 +152,30 @@ def retry_call(f, fargs=None, fkwargs=None, exceptions=Exception, tries=-1, dela args = fargs if fargs else list() kwargs = fkwargs if fkwargs else dict() return __retry_internal(partial(f, *args, **kwargs), exceptions, tries, delay, max_delay, backoff, jitter, logger) + + +def retry_generator(f, fargs=None, fkwargs=None, exceptions=Exception, tries=-1, delay=0, max_delay=None, backoff=1, jitter=0, + logger=logging_logger): + """ + Return a generator that wraps a generator returned by a given function, retrying from the start on failure. + + :param f: the function that returns a generator + :param fargs: the positional arguments of the function to execute. + :param fkwargs: the named arguments of the function to execute. + :param exceptions: an exception or a tuple of exceptions to catch. default: Exception. + :param tries: the maximum number of attempts. default: -1 (infinite). + :param delay: initial delay between attempts. default: 0. + :param max_delay: the maximum value of delay. default: None (no limit). + :param backoff: multiplier applied to delay between attempts. default: 1 (no backoff). + :param jitter: extra seconds added to delay between attempts. default: 0. + fixed if a number, random if a range tuple (min, max) + :param logger: logger.warning(fmt, error, delay) will be called on failed attempts. + default: retry.logging_logger. if None, logging is disabled. + :returns: the result of the f function. + """ + args = fargs if fargs else list() + kwargs = fkwargs if fkwargs else dict() + return __retry_internal_generator(partial(f, *args, **kwargs), exceptions, tries, delay, max_delay, backoff, jitter, logger) + + + diff --git a/tests/test_retry.py b/tests/test_retry.py index 64f45cd..127b239 100644 --- a/tests/test_retry.py +++ b/tests/test_retry.py @@ -14,6 +14,7 @@ from retry.api import retry_call from retry.api import retry +from retry.api import retry_generator def test_retry(monkeypatch): @@ -183,3 +184,162 @@ def f(value=0): assert result == kwargs['value'] assert f_mock.call_count == 1 + + +def test_retry_with_generator(): + """Tests that 3 tries occur and generated result: + + + [1, 2] + + [1, 2, 4] + """ + _vals = [RuntimeError(0), 1, 2, RuntimeError(3), 4] + f_calls = list() + + @retry(exceptions=RuntimeError, generator=True) + def f_mock(): + f_calls.append(()) + for v in _vals: + if isinstance(v, BaseException): + _vals.remove(v) + raise v + else: + yield v + + tries = 3 + expected = [1, 2, 1, 2, 4] + actual = [] + try: + actual.extend(x for x in f_mock()) + except RuntimeError: + pass + + assert len(f_calls) == tries + assert actual == expected + + +def mock_generator(*vals): + _vals = list(vals) + calls = list() + + def tmp_generator(*args, **kwargs): + calls.append((args, kwargs)) + for v in _vals: + if isinstance(v, BaseException): + _vals.remove(v) + raise v + else: + yield v + + return tmp_generator, calls + + +def test_retry_generator(): + """Tests that 3 tries occur and generated result: + + + + """ + f_mock, f_calls = mock_generator(RuntimeError(0), RuntimeError(1)) + tries = 3 + expected = [] + actual = [] + try: + actual.extend(x for x in retry_generator(f_mock, exceptions=RuntimeError, tries=tries)) + except RuntimeError: + pass + + assert len(f_calls) == 3 + assert actual == expected + + +def test_retry_generator_2(): + """Tests that 3 tries occur and generated result: + + + + [2] + """ + f_mock, f_calls = mock_generator(RuntimeError(0), RuntimeError(1), 2) + tries = 5 + expected = [2] + actual = [] + try: + actual.extend(x for x in retry_generator(f_mock, exceptions=RuntimeError, tries=tries)) + except RuntimeError: + pass + + assert len(f_calls) == 3 + assert actual == expected + + +def test_retry_generator_3(): + """Tests that 3 tries occur and generated result: + + [0] + + [0] + + [0] + """ + f_mock, f_calls = mock_generator(0, RuntimeError(1), RuntimeError(2)) + tries = 5 + expected = [0, 0, 0] + actual = [] + try: + actual.extend(x for x in retry_generator(f_mock, exceptions=RuntimeError, tries=tries)) + except RuntimeError: + pass + + assert len(f_calls) == 3 + assert actual == expected + + +def test_retry_generator_4(): + """Tests that 2 tries (vs. 3 for test above, due to parameter tries) occur and generated result: + + [0] + + [0] + + """ + f_mock, f_calls = mock_generator(0, RuntimeError(1), RuntimeError(2)) + tries = 2 + expected = [0, 0] + actual = [] + try: + actual.extend(x for x in retry_generator(f_mock, exceptions=RuntimeError, tries=tries)) + except RuntimeError: + pass + + assert len(f_calls) == 2 + assert actual == expected + + +def test_retry_generator_with_args(): + """Tests args are used""" + f_mock, f_calls = mock_generator('test') + tries = 2 + args = (1, 3) + try: + _ = [x for x in retry_generator(f_mock, fargs=args, exceptions=RuntimeError, tries=tries)] + except RuntimeError: + pass + + assert len(f_calls) == 1 + assert f_calls[0][0] == args + + +def test_retry_generator_with_kwargs(): + """Tests kwargs are used""" + f_mock, f_calls = mock_generator('test') + tries = 2 + kwargs = {'a': 1, 'b': 3} + try: + _ = [x for x in retry_generator(f_mock, fkwargs=kwargs, exceptions=RuntimeError, tries=tries)] + except RuntimeError: + pass + + assert len(f_calls) == 1 + assert f_calls[0][1] == kwargs