怎么实现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():
    # 模拟一个可能会失败的操作
    import random
    if random.random() < 0.5:
        raise ValueError("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):
    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(3)
def long_running_task():
    import time
    time.sleep(5)
    return "Task completed"

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

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

4. 结合重试和超时功能

最后，我们将重试和超时功能结合起来，实现一个既能重试又能超时的装饰器。

def retry_timeout(max_retries=3, delay=1, timeout_seconds=5):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            retries = 0
            while retries < max_retries:
                try:
                    return timeout(timeout_seconds)(func)(*args, **kwargs)
                except (TimeoutError, Exception) 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=3)
def fetch_data_with_timeout():
    import random
    import time
    if random.random() < 0.5:
        time.sleep(4)  # 模拟超时
    return "Data fetched successfully"

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

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

5. 总结

通过使用装饰器，我们可以轻松地为函数添加重试和超时功能，而无需修改原有函数的代码。这种方法不仅提高了代码的可读性，还增强了代码的健壮性。希望本文对你理解和使用Python装饰器有所帮助。