如何实现Python重试超时装饰器

在编写Python程序时，我们经常会遇到需要重试某些操作的情况，尤其是在处理网络请求、数据库连接等可能因外部因素失败的操作时。为了简化代码并提高可读性，我们可以使用装饰器来实现重试和超时功能。本文将介绍如何实现一个Python重试超时装饰器。

1. 什么是装饰器？

装饰器是Python中的一种高级函数，它允许我们修改或扩展其他函数的行为，而不需要修改这些函数的源代码。装饰器通常用于日志记录、权限检查、性能测试等场景。

2. 实现重试装饰器

首先，我们来实现一个简单的重试装饰器。这个装饰器会在函数执行失败时自动重试指定的次数。

import time
from functools import wraps

def retry(max_retries=3, delay=1):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            retries = 0
            while retries < max_retries:
                try:
                    return func(*args, **kwargs)
                except Exception as e:
                    retries += 1
                    if retries == max_retries:
                        raise e
                    time.sleep(delay)
        return wrapper
    return decorator

使用示例

@retry(max_retries=5, delay=2)
def fetch_data():
    # 模拟一个可能失败的操作
    if random.random() < 0.5:
        raise Exception("Failed to fetch data")
    return "Data fetched successfully"

print(fetch_data())

在这个例子中，fetch_data函数会在失败时最多重试5次，每次重试之间等待2秒。

3. 实现超时装饰器

接下来，我们来实现一个超时装饰器。这个装饰器会在函数执行时间超过指定时间时抛出超时异常。

import signal
from functools import wraps

class TimeoutError(Exception):
    pass

def timeout(seconds=10):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            def handle_timeout(signum, frame):
                raise TimeoutError("Function timed out")

            signal.signal(signal.SIGALRM, handle_timeout)
            signal.alarm(seconds)
            try:
                result = func(*args, **kwargs)
            finally:
                signal.alarm(0)
            return result
        return wrapper
    return decorator

使用示例

@timeout(seconds=5)
def long_running_task():
    time.sleep(10)
    return "Task completed"

try:
    print(long_running_task())
except TimeoutError as e:
    print(e)

在这个例子中，long_running_task函数如果执行时间超过5秒，将会抛出TimeoutError异常。

4. 结合重试和超时装饰器

我们可以将重试和超时装饰器结合起来，以实现一个既支持重试又支持超时的装饰器。

def retry_timeout(max_retries=3, delay=1, timeout_seconds=10):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            retries = 0
            while retries < max_retries:
                try:
                    return timeout(timeout_seconds)(func)(*args, **kwargs)
                except (Exception, TimeoutError) as e:
                    retries += 1
                    if retries == max_retries:
                        raise e
                    time.sleep(delay)
        return wrapper
    return decorator

使用示例

@retry_timeout(max_retries=5, delay=2, timeout_seconds=5)
def fetch_data_with_timeout():
    # 模拟一个可能失败的操作
    if random.random() < 0.5:
        raise Exception("Failed to fetch data")
    time.sleep(6)  # 模拟超时
    return "Data fetched successfully"

try:
    print(fetch_data_with_timeout())
except Exception as e:
    print(e)

在这个例子中，fetch_data_with_timeout函数会在失败时最多重试5次，每次重试之间等待2秒，并且每次执行时间超过5秒时会抛出超时异常。

5. 总结

通过使用装饰器，我们可以轻松地为函数添加重试和超时功能，而无需修改函数的原始代码。这不仅提高了代码的可读性，还增强了程序的健壮性。希望本文能帮助你理解如何实现Python中的重试超时装饰器，并在实际项目中应用这些技巧。