Unconventional Ways to Make Python Retry Code More Elegant

Loop Retry

A complete retry process roughly includes loop retries, exception handling, delay waiting, and subsequent operations. The standard way to write this is using a for loop with try ... except ... to catch exceptions. A reference code is as follows:

import time
from random import random

allright = False  # Success flag

for i in range(5):  # Retry up to 5 times
    try:
        # Code that might fail
        x = random()
        if x < 0.5:
            yyyy  # yyyy is undefined, so it will throw an error
        allright = True
        break
    except Exception as e:
        print(e)  # Print error message
        if i < 4:
            time.sleep(2)  # Delay for two seconds

if allright:
    # Perform some operations
    print('Execution successful')
else:
    # Perform other operations
    print('Execution failed')

Our goal is to simplify the code before if allright:. We can see that it follows a fixed format: a for loop combined with a template of "try ... break ... except ... sleep ...". It is easy to imagine that there is significant room for simplification.

Function Decoration

The problem with the for loop is that if there are many places requiring retries and the exception handling logic is the same, rewriting the except code every time becomes redundant. In this case, the standard recommended method is to wrap the error-prone code into a function and write a decorator to handle exceptions:

import time
from random import random

def retry(f):
    """Retry decorator, adds retry functionality to the wrapped function
    """
    def new_f(*args, **kwargs):
        for i in range(5):  # Retry up to 5 times
            try:
                return True, f(*args, **kwargs)
            except Exception as e:
                print(e)  # Print error message
                if i < 4:
                    time.sleep(2)  # Delay for two seconds
        return False, None
    return new_f

@retry
def f():
    # Code that might fail
    x = random()
    if x < 0.5:
        yyyy  # yyyy is undefined, so it will throw an error
    return x

allright, _ = f()  # Returns execution status and result
if allright:
    # Perform some operations
    print('Execution successful')
else:
    # Perform other operations
    print('Execution failed')

When multiple different code blocks need retries, you only need to write them as functions and add the @retry decorator to implement the same retry logic. Thus, the decorator approach is indeed a concise and intuitive solution, which explains why it has become the standard. Most mainstream retry libraries, such as tenacity, or earlier ones like retry and retrying, are based on the decorator principle.

Ideal Syntax

However, while the decorator approach is standard, it is not perfect. First, encapsulating the retry code into a separate function can make the code flow feel interrupted or "stuttery" in many cases. Second, because the code is encapsulated in a function, intermediate variables within the scope cannot be used directly; any variables needed must be passed through arguments or returned, which feels a bit roundabout. Overall, while decorators simplify retry code, they still leave something to be desired.

The perfect retry code I imagine should be based on a context manager, something like:

with Retry(max_tries=5) as retry:
    # Code that might fail
    x = random()
    if x < 0.5:
        yyyy  # yyyy is undefined, so it will throw an error

if retry.allright:
    # Perform some operations
    print('Execution successful')
else:
    # Perform other operations
    print('Execution failed')

However, after studying the principles of context managers, I realized that this ideal syntax is destined to be impossible to achieve. This is because a context manager can only manage the "context" (the setup and teardown) but cannot manage the "body" of the code (the "code that might fail" in this article). Specifically, a context manager is a class with __enter__ and __exit__ methods. It inserts __enter__ before the code runs and __exit__ after it runs, but it cannot control the code in the middle (e.g., making it run multiple times).

Therefore, the attempt to implement a one-line retry based solely on a context manager has failed.

A Bit of a Struggle

The good news, however, is that while a context manager cannot implement a loop, its __exit__ method can handle exceptions. So, it can at least replace the "try ... except ..." block to handle exceptions. Thus, we can write:

import time
from random import random

class Retry:
    """Custom context manager for handling exceptions
    """
    def __enter__(self):
        self.allright = False
        return self
    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is None:
            self.allright = True
        else:
            print(exc_val)
            time.sleep(2)
        return True

for i in range(5):  # Retry up to 5 times
    with Retry() as retry:
        # Code that might fail
        x = random()
        if x < 0.5:
            yyyy  # yyyy is undefined, so it will throw an error
        break

if retry.allright:
    # Perform some operations
    print('Execution successful')
else:
    # Perform other operations
    print('Execution failed')

This latest version is actually very close to the ideal syntax mentioned in the previous section. There are two differences: 1. You need to write an extra for loop, but this is unavoidable because, as stated earlier, a context manager cannot initiate a loop, so you must use for or while externally; 2. You need to explicitly add a break, which we can find a way to optimize away.

Additionally, this version has a small flaw: if all retries fail, the code will still sleep after the final failure. Theoretically, this is unnecessary and should be removed.

Further Optimization

To optimize away the break, the loop needs to learn to stop itself. There are two ways to do this: the first is to use a while loop and change the stopping condition based on the retry result, which leads to a result similar to "Handling exceptions inside context managers". The second method is to keep the for loop but replace range(5) with an iterator that changes based on the retry result. This article explores the latter.

After analysis, I found that by using the built-in methods __call__ and __iter__, retry can act as a mutable iterator while also solving the unnecessary sleep issue after the final failure:

import time
from random import random

class Retry:
    """Context manager + Iterator for handling exceptions
    """
    def __call__(self, max_tries=5):
        self.max_tries = max_tries
        return self
    def __iter__(self):
        for i in range(self.max_tries):
            yield i
            if self.allright or i == self.max_tries - 1:
                return
            time.sleep(2)
    def __enter__(self):
        self.allright = False
    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is None:
            self.allright = True
        else:
            print(exc_val)
        return True

retry = Retry()
for i in retry(5):  # Retry up to 5 times
    with retry:
        # Code that might fail
        x = random()
        if x < 0.5:
            yyyy  # yyyy is undefined, so it will throw an error

if retry.allright:
    # Perform some operations
    print('Execution successful')
else:
    # Perform other operations
    print('Execution failed')

Attentive readers might wonder: you managed to remove one break, but added retry = Retry(). The total number of lines remains the same (and the context manager is more complex). Is it worth the trouble? In fact, the retry object here is reusable. A user only needs to define retry = Retry() once, and then in subsequent retries, they only need:

for i in retry(max_tries):
    with retry:
        # Code that might fail

This is already extremely close to the ideal implementation.

Ultimate Version

However, "define retry = Retry() once and reuse it" is only suitable for single-process environments. In multi-processing, you would still need to define retry = Retry() separately. Furthermore, this kind of reuse gives a feeling that "different retries are not completely isolated." Is it possible to remove that line entirely? After some thought, I found it is possible! Here is the reference code:

import time
from random import random

class Retry:
    """Context manager + Iterator for handling exceptions
    """
    def __init__(self, max_tries=5):
        self.max_tries = max_tries
    def __iter__(self):
        for i in range(self.max_tries):
            yield self
            if self.allright or i == self.max_tries - 1:
                return
            time.sleep(2)
    def __enter__(self):
        self.allright = False
    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is None:
            self.allright = True
        else:
            print(exc_val)
        return True

for retry in Retry(5):  # Retry up to 5 times
    with retry:
        # Code that might fail
        x = random()
        if x < 0.5:
            yyyy  # yyyy is undefined, so it will throw an error

if retry.allright:
    # Perform some operations
    print('Execution successful')
else:
    # Perform other operations
    print('Execution failed')

The change this time was to replace __call__ with __init__, and change yield i in __iter__ to yield self, returning the object itself. This way, you don’t need a separate line to initialize retry = Retry(5). Instead, for retry in Retry(5): handles both initialization and alias assignment. Since it re-initializes every time, it achieves complete isolation between retries, killing two birds with one stone.

Summary

This article has explored the implementation of retry mechanisms in Python quite thoroughly, attempting to reach what I consider a perfect implementation. The final result has more or less met my expectations.

However, I must admit that the motivation for this article was essentially a result of "OCD" (obsessive-compulsive disorder) regarding code aesthetics. There is no substantial improvement in algorithmic efficiency. Spending too much time on programming details like this is, to some extent, an "unconventional" or "unproductive" pursuit, and perhaps not something strictly worth emulating.

When reposting, please include the original address: https://kexue.fm/archives/9938

For more detailed reposting matters, please refer to: "Scientific Space FAQ"