Retry function call with context manager

Goal

A function call may fail due to a temporal issue, for instance, a request call to an external URL. In this case, it may be useful, to retry calling that function multiple times in sequence, and only raising an error if all retries fail.

The retry module

The retry package provides a convenient method retry_call, which calls a function either up until the amount of retries is reached or to be retried function returns a non-error type object.

from retry.api import retry_call

def say_hi():
    print("hi")
    raise OSError


def retry_say_hi():
    exceptions = (OSError, IOError,)
    try:
        retry_call(say_hi, tries=3)
    except exceptions:
        print("That didn't go well")


if __name__ == "__main__":
    retry_say_hi()

hi
hi
hi
That didn't go well

The above function call of retry_say_hi() enters say_hi() three times, each time raising a OSError, which causes a retry, and ultimately fails.

Wrapping retry in a context manager

To provide a uniform way to retry sensitive functions, a [context manager] can be used to wrap a to-be-retried-in-error-case function:

from typing import Tuple, Type
from logging import Logger
from contextlib import contextmanager
from dataclasses import dataclass
from functools import partial


@dataclass
class RetryContextManager:
    """Context manager to *retry* calls.

    The class wraps the methods provided by the `retry` Python package
    to simplify the retry of failing function calls. The class implements
    a context manager that allows us to specify retry parameters once
    at init time.
    """

    exceptions: Tuple[Type[Exception], ...] = (Exception,)
    tries: int = 1
    delay: int = 0
    max_delay: int = None
    backoff: int = 1
    jitter: Tuple[int, int] = (0, 0)
    logger: Logger = None

    @contextmanager
    def retry(self, f, *fargs, **fkwargs):
        yield partial(
            retry_call, f, fargs, fkwargs,
            exceptions=self.exceptions,
            tries=self.tries,
            delay=self.delay,
            max_delay=self.max_delay,
            backoff=self.backoff,
            jitter=self.jitter,
            logger=self.logger
        )

Line 9-24: Definition of the RetryContextManager dataclass

Line 26-36: Definition of the context manager factory function. The @contextmanager decorator enables to define a function as context manager, instead of a class implementing the __enter__() and __exit__() method. It returns the retry_call method as partial object, which behaves like the regular method would, but passes in the arguments from the RetryContextManager instance it is called from. The partial method allows calling the function (here: retry_call) with custom positional arguments.

Using the custom context manager

A class may now define a custom RetryContextManager, specifying custom tries or delay parameters

import logging
import requests

class MyRequestClass:

    def __init__(self):
        self.get_response_retry_ctx = RetryContextManager(
            tries=3,
            delay=1,
            exceptions=(requests.RequestException,),
            logger=logging.getLogger("RobotFramework")
        )

    def get_audio_bytes_from_url(self, url: str) -> bytes:
        with self.get_response_retry_ctx.retry(requests.get,
                                               url) as get_response:
            response = get_response()
        if response.status_code == 200:
            print("That went well")
        else:
            raise ConnectionError(f"Unable to receive content from {url}")

Line 7: An instance of RetryContextManager dataclass which contains the retry context manager method is defined.

Line 15 & 16: The retry context manager is entered, passing in requests.get as the function to be called and url as a function argument (which leads to the function call requests.get(url)) an assigning the function call to get_response.

Line 17: The get_response function is called. Afterward, the context manager is closed.