The Sidetrack of Making Python's Retry Code More Elegant
In this post, we'll dig into a programming puzzle: how to implement retries more elegantly in Python.
In the post Happy New Year! Notes on the Development Experience of Cool Papers, I shared some experiences from developing Cool Papers, including some of the network communication steps that Cool Papers relies on. Whenever network communication is involved, there's always a risk of failure (nobody can guarantee the network won't intermittently misbehave), so retrying is a basic operation for network communication. Beyond that, retry mechanisms are also typically needed whenever multiprocessing, databases, or hardware interactions are involved.
Implementing retries in Python isn't hard, but doing so in a way that's both simple and readable takes a bit of skill. What follows is a record of my own attempts at this.
Retrying with a loop
A complete retry flow generally consists of a retry loop, exception handling, a delay/wait, and follow-up actions. The standard way to write this is with a for loop, using "try ... except ..." to catch exceptions. Here's a reference implementation: more
import time
from random import random
allright = False # 执行成功的标记
for i in range(5): # 最多重试5次
try:
# 有概率出错的代码
x = random()
if x < 0.5:
yyyy # 未定义yyyy，所以会报错
allright = True
break
except Exception as e:
print(e) # 打印错误信息
if i < 4:
time.sleep(2) # 延时两秒
if allright:
# 执行某些操作
print('执行成功')
else:
# 执行另一些操作
print('执行失败')
Our goal from here on is to simplify the code before if allright:. As you can see, it follows a fairly fixed pattern: a for loop wrapped around the template "try ... break ... except ... sleep ...". It's not hard to imagine that there's plenty of room to simplify this.
Function decorators
The problem with the for-loop approach is that if there are many places in the code that need retrying, and the exception-handling logic is the same each time, then rewriting the except code over and over becomes tedious. In this situation, the standard recommendation is to wrap the error-prone code into a function, and write a decorator to handle the exceptions:
import time
from random import random
def retry(f):
"""重试装饰器，包装函数加上重试功能
"""
def new_f(*args, **kwargs):
for i in range(5): # 最多重试5次
try:
return True, f(*args, **kwargs)
except Exception as e:
print(e) # 打印错误信息
if i < 4:
time.sleep(2) # 延时两秒
return False, None
return new_f
@retry
def f():
# 有概率出错的代码
x = random()
if x < 0.5:
yyyy # 未定义yyyy，所以会报错
return x
allright, _ = f() # 返回执行状态和执行结果
if allright:
# 执行某些操作
print('执行成功')
else:
# 执行另一些操作
print('执行失败')
When multiple different pieces of code all need retrying, you just need to turn each of them into a function and slap the @retry decorator on top to get the same retry logic. So the decorator-based approach really is a clean solution, and quite intuitive too — it's easy to see why it became the standard. The mainstream retry libraries today, such as tenacity, or the older retry and retrying, are all built on this decorator principle.
The ideal way to write it
That said, while the decorator approach is standard, it isn't perfect. First, you need to wrap the retryable code into a separate function, which in many cases breaks the flow of the code, giving it a jarring, "stop-and-go" feel. Second, because the code is wrapped inside a function, any intermediate variables inside it aren't directly accessible — anything you need has to be passed in through return, which feels like an awkward workaround. All in all, while decorators do simplify retry code, something is still missing.
The perfect retry code I have in mind should be based on a context manager, something like:
with Retry(max_tries=5) as retry:
# 有概率出错的代码
x = random()
if x < 0.5:
yyyy # 未定义yyyy，所以会报错
if retry.allright:
# 执行某些操作
print('执行成功')
else:
# 执行另一些操作
print('执行失败')
However, after digging into how context managers actually work, I realized that this ideal form is fundamentally impossible to achieve. A context manager can only manage the surrounding context — it cannot control the main body of code (i.e., the "error-prone code" we've been discussing). Specifically, a context manager is a class with __enter__ and __exit__ methods, which insert __enter__ before the code runs (the "before" part) and __exit__ after the code runs (the "after" part), but it has no way to manipulate the code in between (e.g., making it run multiple times).
So, the dream of implementing retries in a single line using a context manager is dead on arrival.
A bit of a struggle
The good news, though, is that although a context manager can't implement looping, its __exit__ method can handle exceptions. So it can at least take the place of "try ... except ...". This gets us to:
import time
from random import random
class Retry:
"""自定义处理异常的上下文管理器
"""
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): # 最多重试5次
with Retry() as retry:
# 有概率出错的代码
x = random()
if x < 0.5:
yyyy # 未定义yyyy，所以会报错
break
if retry.allright:
# 执行某些操作
print('执行成功')
else:
# 执行另一些操作
print('执行失败')
This latest version is actually already very close to the ideal form from the previous section. There are two differences: first, we still need an extra for loop line — this is unavoidable, since, as already discussed, a context manager can't drive a loop on its own, so we have to start the loop ourselves with either for or while. Second, we need to explicitly add a line for break, which we might be able to find a way to optimize away.
There's also a small flaw in this version: if every retry attempt fails, then after the very last failed attempt, it'll still call sleep, which is unnecessary in principle and should ideally be eliminated.
Further optimization
To eliminate break, the loop needs to learn to stop on its own, and there are two ways to do this. The first is to switch to a while loop and change the stopping condition depending on the retry outcome — this leads to a result similar to Handling exceptions inside context managers. The second is to keep the for loop, but replace range(5) with an iterator whose behavior changes based on the retry outcome. This post focuses on the latter approach.
After some analysis, I found that using the built-in methods __call__ and __iter__, we can make retry double as a mutable iterator, and also solve the problem of the unnecessary sleep after the final failed attempt:
import time
from random import random
class Retry:
"""处理异常的上下文管理器 + 迭代器
"""
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): # 最多重试5次
with retry:
# 有概率出错的代码
x = random()
if x < 0.5:
yyyy # 未定义yyyy，所以会报错
if retry.allright:
# 执行某些操作
print('执行成功')
else:
# 执行另一些操作
print('执行失败')
Careful readers might object: you found a way to remove one line, break, but added another, retry = Retry(), so the total line count hasn't changed (and you've made the context manager more complicated to boot) — was this really worth the trouble? In fact, the retry object here is reusable: the user only needs to define retry = Retry() once, and then for every subsequent retry, all that's needed is:
for i in retry(max_tries):
with retry:
# 有概率出错的代码
That's it. So although the context manager itself has gotten a bit more complex, this is now an implementation that's about as close to the ideal as we can get.
The ultimate version
That said, "define retry = Retry() once, reuse retry many times" only works well within a single process — for multiprocessing, you'd still need to define retry = Retry() separately for each process. There's also something slightly unsatisfying about this kind of reuse: it gives the impression that different retry sequences aren't fully isolated from one another. Is there a way to eliminate this line entirely? After some more thought, I found that yes, it is possible! Here's the reference code:
import time
from random import random
class Retry:
"""处理异常的上下文管理器 + 迭代器
"""
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): # 最多重试5次
with retry:
# 有概率出错的代码
x = random()
if x < 0.5:
yyyy # 未定义yyyy，所以会报错
if retry.allright:
# 执行某些操作
print('执行成功')
else:
# 执行另一些操作
print('执行失败')
The change this time is swapping __call__ for __init__, and then changing yield i inside __iter__ to yield self, i.e., having it return the object itself. This way, there's no need for a separate initialization line retry = Retry(5) — instead, for retry in Retry(5): handles both initialization and the alias assignment at once. And since every retry attempt gets freshly initialized, this also achieves full isolation between separate retry sequences — killing two birds with one stone.
Summary
This post has taken a fairly thorough look at how to write retry logic in Python, in an attempt to arrive at what I consider to be the ideal implementation of retry code. In the end, the result more or less lives up to what I had in mind.
That said, I have to admit that the real motivation behind this post was essentially just my own "OCD" acting up — there's no substantive improvement in algorithmic efficiency here. Spending this much time on programming minutiae is, in some sense, a bit of a sidetrack, a distraction from more serious work, and not something particularly worth emulating.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.