retry.py

from functools import partial
from itertools import imap
from operator import or_

def retry(ntimes=3, ignore=None, trap=None):
    """Retry an operation some number of times.

    The ignore and trap arguments may be a sequence (or a single)
    exception class. If the decorated function raises an exception
    matching ignore, it will be raised.

    If the decorated function raises an exception matching trap, the
    function will be retried up to ntimes. If it continues to raise
    exceptions, the last exception is raised.

    Use:

    # Retry any exception up to three times
    @retry(3)
    def foo():
        return flaky_operation()

    # Retry any exception (except ImportantException) once
    @retry(1, ImportantException)
    def foo():
        return flaky_operation()
    """

    ignore = ((ignore,) if isinstance(ignore, Exception) else ignore) or ()
    trap = ((trap,) if isinstance(trap, Exception) else trap) or (Exception,)

    instance = lambda obj, types=None: \
        reduce(or_, imap(partial(isinstance, obj), types or ()), False)

    def __closure__(func):
        def __inner__(*args, **kwargs):
            attempt = 0
            while attempt < ntimes:
                try:
                    return func(*args, **kwargs)
                except Exception, e:
                    if instance(e, ignore) or not instance(e, trap):
                        raise e
                    attempt += 1
            raise e
        return __inner__
    return __closure__