retry context manager

retry.py

import time

from logging import getLogger


class RetryContextInvalidResult(Exception):
    pass


class RetryContext(object):
    """
    Usage:

    with retry_request(max_retries=3, delay=0.1) as manager:
        response = manager.execute(requests.get, 'http://google.com')
    """
    def __init__(self, max_retries, validator=lambda r: True, delay=None):
        self.max_retries = max_retries
        self.validator = validator
        self.delay = delay
        self.logger = getLogger('utils.requests.RetryContext')

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, traceback):
        return False

    def execute(self, method, *args, **kwargs):
        error = None
        for _ in range(self.max_retries):
            try:
                result = method(*args, **kwargs)
            except Exception as error:
                self.logger.warning(error.message)
            else:
                if self.validator(result):
                    return result
                error = RetryContextInvalidResult(
                    'Invalid result from <{method_module}.{method_name}'
                    ' args={args}, kwargs={kwargs}> ...failed validator:'
                    ' {validator_module}.{validator_name}'.format(
                        method_module=method.__module__,
                        method_name=method.__name__,
                        args=args,
                        kwargs=kwargs,
                        validator_module=self.validator.__module__,
                        validator_name=self.validator.__name__,
                    )
                )
                # if we run out of retries the exception will be re-raised,
                # so don't confuse the logs with an ERROR message here:
                self.logger.warning(error.message)

            if self.delay:
                time.sleep(self.delay)

        raise error


def retry_request(max_retries, delay=1):
    '''
    We want to retry in the case of transient http connection error, i.e.
    where no response would be returned.

    In case of successful response, http error statuses should be handled
    in your own code.
    '''
    return RetryContext(
        max_retries=max_retries,
        delay=delay,
    )