retry generator + tests

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:
+        # <Initial call>
+        # <New call due to RuntimeError(0)>
+        # 1, 2
+        # <New call due to RuntimeError(3)>
+        # 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.

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,15 +109,22 @@ 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.
     """
 
     @decorator
     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)
+
+
+

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:
+    <Initial call>
+    <New call due to RuntimeError(0)>
+    [1, 2]
+    <New call due to RuntimeError(3)>
+    [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:
+    <Initial call>
+    <New call due to RuntimeError(0)>
+    <New call due to RuntimeError(1)>
+    """
+    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:
+    <Initial call>
+    <New call due to RuntimeError(0)>
+    <New call due to RuntimeError(1)>
+    [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:
+    <Initial call>
+    [0]
+    <New call due to RuntimeError(1)>
+    [0]
+    <New call due to RuntimeError(2)>
+    [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:
+    <Initial call>
+    [0]
+    <New call due to RuntimeError(1)>
+    [0]
+    <New call due to RuntimeError(2)>
+    """
+    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