Python装饰器之时间装饰器

一、需求引入

在日常工作中，经常会需要对一些方法的执行耗时进行统计，以方便优化性能；在一些自动化测试时需要判断被测对象的执行耗时是否超时。要实现这些功能的，并且可复用的话，装饰器是一个不错的选择。

二、计算执行耗时装饰器

同步方法装饰器

import time


def timeit(func):
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        end_time = time.time()

        print(f"{func.__name__} finished in {end_time - start_time:.6f} seconds.")

        return result

    return wrapper


# 示例同步函数
@timeit
def some_sync_task(sleep_duration: float):
    time.sleep(sleep_duration)


# 运行示例
if __name__ == "__main__":
    some_sync_task(0.5) # 让同步任务模拟运行2秒

异步方法装饰器

import asyncio
import time


def async_timeit(func):
    async def wrapper(*args, **kwargs):
        start_time = time.time()
        await func(*args, **kwargs)
        end_time = time.time()

        print(f"{func.__name__} finished in {end_time - start_time:.6f} seconds.")

    return wrapper


# 示例异步函数
@async_timeit
async def some_async_task(sleep_duration: float):
    await asyncio.sleep(sleep_duration)


# 运行示例
async def main():
    await some_async_task(0.5)  # 让异步任务模拟运行2秒


if __name__ == "__main__":
    asyncio.run(main())

三、超时装饰器

其实我一开始使用的timeout参数是int类型的，但考虑到超时不一定是整数值，比如0.5秒之类的是现实中更实用的场景，因此把timeout参数改为float类型。

1. 丐版超时装饰器

适用于对执行耗时比较敏感，需要尽量减少装饰器本身耗时的场景。

缺点是比较简陋，只能以秒为单位设置timeout.

# 异步装饰器（包含实际执行耗时）
import asyncio
import time
from functools import wraps

class TimeOutErr(Exception):
    def __init__(self, message: str = "Function execution timed out", exec_time: float = None):
        self.message = message
        self.exec_time = exec_time
        super().__init__(message)


def async_timeout(timeout: float):
    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            start_time = time.time()

            try:
                result = await asyncio.wait_for(func(*args, **kwargs), timeout=timeout)

                return result

            except asyncio.TimeoutError:
                end_time = time.time()
                raise TimeOutErr(
                    f"Function '{func.__name__}' exceeded the timeout of {timeout} seconds after running for {end_time - start_time:.6f} seconds.")

        return wrapper

    return decorator


# 同步装饰器（包含实际执行耗时）
def timeout(timeout: float):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            start_time = time.time()

            try:
                result = func(*args, **kwargs)

                if time.time() - start_time > timeout:
                    end_time = time.time()
                    raise TimeOutErr(
                        f"Function '{func.__name__}' exceeded the timeout of {timeout} seconds after running for {end_time - start_time:.6f} seconds.")

                return result

            except Exception as e:
                if isinstance(e, TimeOutErr):
                    raise e
                else:
                    raise

        return wrapper

    return decorator


@async_timeout(2.0)
async def some_async_task(sleep_duration: float):
    await asyncio.sleep(sleep_duration)
    return "Execution completed."



@timeout(2.0)
def some_sync_task(sleep_duration: float):
    time.sleep(sleep_duration)
    return "Execution completed."


# 运行示例
async def main():
    try:
        print("Async Task:")
        result = await some_async_task(1.5)  # 执行成功
        print(f"Task result: {result}")

        try:
            result = await some_async_task(3)  # 将抛出TimeOutErr异常
            print(f"Task result: {result}")
        except TimeOutErr as e:
            print(f"Caught exception: {e}")

    except Exception as e:
        print(f"Unexpected error in async task: {e}")

    try:
        print("\nSync Task:")
        result = some_sync_task(1.5)  # 执行成功
        print(f"Task result: {result}")

        try:
            result = some_sync_task(3)  # 将抛出TimeOutErr异常
            print(f"Task result: {result}")
        except TimeOutErr as e:
            print(f"Caught exception: {e}")

    except Exception as e:
        print(f"Unexpected error in sync task: {e}")


if __name__ == "__main__":
    asyncio.run(main())

2. 支持不同时间单位的超时装饰器

允许用户通过minutes、seconds等命名参数来指定超时时间。

你可以在装饰器中直接通过命名参数设置超时时间，例如@async_timeout(minutes=1, seconds=30)。

请注意，这里我们假设如果设置了多个单位，则它们会累加计算总的超时时间。

在函数执行完毕后检查是否超过了设定的超时时间，并根据需要抛出TimeOutErr异常。

同步装饰器

class TimeOutErr(Exception):
    def __init__(self, message, exec_time=None):
        self.message = message
        self.exec_time = exec_time
        super().__init__(message)


def sync_timeout(**kwargs):
    # 默认超时时间为0秒
    timeout_seconds = 0.0

    if 'minutes' in kwargs:
        timeout_seconds += kwargs['minutes'] * 60.0
    if 'seconds' in kwargs:
        timeout_seconds += kwargs['seconds']
    if 'milliseconds' in kwargs:  # 支持毫秒
        timeout_seconds += kwargs['milliseconds'] / 1000.0
    if 'microseconds' in kwargs:
        timeout_seconds += kwargs['microseconds'] / 1000000.0

    if timeout_seconds <= 0:
        raise ValueError("Timeout value must be greater than zero.")

    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs2):
            start_time = time.time()

            try:
                result = func(*args, **kwargs2)

                end_time = time.time()
                if end_time - start_time > timeout_seconds:
                    raise TimeOutErr(
                        f"Function '{func.__name__}' exceeded the timeout after running for {end_time - start_time:.6f} seconds.")

                return result

            except Exception as e:
                if isinstance(e, TimeOutErr):
                    raise e
                else:
                    raise

        return wrapper

    return decorator


# 示例代码
@sync_timeout(minutes=1, seconds=30)
def some_sync_task(sleep_duration: float):
    time.sleep(sleep_duration)
    return "Execution completed."


# 运行示例
if __name__ == "__main__":
    try:
        print(some_sync_task(45))  # 将抛出TimeOutErr异常
    except TimeOutErr as e:
        print(f"Caught exception: {e}")

异步装饰器

import asyncio
import time
from typing import Callable, Any, Union, Optional


class TimeOutErr(Exception):
    def __init__(self, message, exec_time=None):
        self.message = message
        self.exec_time = exec_time
        super().__init__(message)


def async_timeout(**kwargs):
    # 默认超时时间为0秒
    timeout_seconds = 0.0

    if 'minutes' in kwargs:
        timeout_seconds += kwargs['minutes'] * 60.0
    if 'seconds' in kwargs:
        timeout_seconds += kwargs['seconds']
    if 'milliseconds' in kwargs:  # 支持毫秒
        timeout_seconds += kwargs['milliseconds'] / 1000.0
    if 'microseconds' in kwargs:
        timeout_seconds += kwargs['microseconds'] / 1000000.0

    if timeout_seconds <= 0:
        raise ValueError("Timeout value must be greater than zero.")

    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs2):
            start_time = time.time()

            try:
                result = await asyncio.wait_for(func(*args, **kwargs2), timeout=timeout_seconds)

                return result

            except asyncio.TimeoutError:
                end_time = time.time()
                raise TimeOutErr(
                    f"Function '{func.__name__}' exceeded the timeout after running for {end_time - start_time:.6f} seconds.")

        return wrapper

    return decorator


# 示例代码
@async_timeout(minutes=1, seconds=30)
async def some_async_task(sleep_duration: float):
    await asyncio.sleep(sleep_duration)
    return "Execution completed."

四、关于装饰器增加耗时的一点思考

在3-2章节中，我们引入了对不同时间单位（如分钟、秒和毫秒）的支持，以提升用户使用的便捷性。然而，随之而来的一个顾虑是这样的改进是否会增加总体的执行耗时。

时间单位支持与性能权衡

增加便利性的考量

增加便利性当然是好的，但也确实存在一种担忧：在装饰器内部进行单位转换和计算可能带来微小的执行开销。

注意：虽然装饰器中的单位处理会占用一定的时间，但与被装饰函数的实际运行时间相比，这部分开销通常是可以忽略不计的。

在大多数实际应用情境下，装饰器初始化及转换所消耗的时间成本远低于整个函数或异步任务本身的执行时间。

可读性与易用性优势

此外，从代码的可读性和开发效率角度来看，提供灵活的超时设置方式能够极大地简化开发流程，并有效减少因单位转换引发的潜在错误。因此，在大部分情况下，这种为了提高易用性而付出的微小代价是完全值得的，特别是当我们的重点在于确保程序逻辑正确且优化整体性能时。

高性能场景下的考量

当然，在高度依赖实时响应或对性能要求极为严苛的环境中，任何额外的性能损耗都需要谨慎评估。在这种情况下，开发者可以选择牺牲一定的便利性来换取几乎可以忽略的性能提升。然而，在常规的应用开发实践中，为了保持代码的整洁与易于维护，采用上述带有时间单位灵活性的装饰器设计方法是可行且推荐的。