如何實現Python重試超時裝飾器

如何實現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中的重試超時裝飾器，并在實際項目中應用這些技巧。