发布于 由admin

`装饰器`的`高级`用法：带参数的装饰器、类装饰器和`wraps`的正确使用。

装饰器的高级用法：带参数的装饰器、类装饰器和 wraps 的正确使用

大家好，今天我们来深入探讨 Python 装饰器的高级用法。装饰器是 Python 中一个非常强大且常用的特性，它可以让我们在不修改原有函数代码的基础上，增加额外的功能。今天，我们将重点讲解带参数的装饰器、类装饰器以及 wraps 的正确使用，帮助大家更好地理解和运用装饰器。

1. 带参数的装饰器

普通的装饰器，比如：

def my_decorator(func):
    def wrapper():
        print("Before the function call.")
        func()
        print("After the function call.")
    return wrapper

@my_decorator
def say_hello():
    print("Hello!")

say_hello()

输出:

Before the function call.
Hello!
After the function call.

这种装饰器只能提供固定的功能增强。如果我们需要根据不同的场景来定制装饰器的行为，就需要使用带参数的装饰器。

1.1 实现原理

带参数的装饰器实际上是返回装饰器的工厂函数。这意味着我们需要编写一个函数，该函数接收参数，并返回一个装饰器函数。这个装饰器函数再接收要装饰的函数，并返回包装后的函数。

1.2 示例代码

def repeat(num_times):
    def decorator_repeat(func):
        def wrapper(*args, **kwargs):
            for _ in range(num_times):
                result = func(*args, **kwargs)
            return result
        return wrapper
    return decorator_repeat

@repeat(num_times=3)
def greet(name):
    print(f"Hello, {name}!")

greet("Alice")

输出:

Hello, Alice!
Hello, Alice!
Hello, Alice!

1.3 代码解释

  • repeat(num_times): 这是一个函数，它接收参数 num_times，并返回一个装饰器 decorator_repeat。
  • decorator_repeat(func): 这是一个装饰器函数，它接收被装饰的函数 func，并返回包装后的函数 wrapper。
  • wrapper(*args, **kwargs): 这是一个包装函数，它执行被装饰的函数 func num_times 次，并将结果返回。

1.4 使用场景

带参数的装饰器在以下场景中非常有用：

  • 日志记录: 可以根据日志级别来决定是否记录日志。
  • 权限控制: 可以根据用户角色来决定是否允许访问某个函数。
  • 缓存: 可以根据缓存过期时间来决定是否从缓存中获取结果。
  • 重试机制: 可以根据重试次数来决定函数失败后重试的次数。

1.5 更复杂的例子 – 缓存装饰器

import functools
import time

def cache(max_age=60):
    def decorator_cache(func):
        cached_value = None
        last_updated = None

        @functools.wraps(func) # 确保保留原始函数的元数据
        def wrapper(*args, **kwargs):
            nonlocal cached_value, last_updated  # 允许修改封闭作用域中的变量

            if cached_value is None or (time.time() - last_updated) > max_age:
                cached_value = func(*args, **kwargs)
                last_updated = time.time()
                print(f"Function executed and value cached. Last updated: {last_updated}")
            else:
                print("Value retrieved from cache.")
            return cached_value
        return wrapper
    return decorator_cache

@cache(max_age=10)
def fetch_data_from_api(api_endpoint):
    print(f"Fetching data from {api_endpoint}...")
    time.sleep(2)  # 模拟耗时操作
    return f"Data from {api_endpoint} at {time.time()}"

print(fetch_data_from_api("https://example.com/api/data"))
time.sleep(5)
print(fetch_data_from_api("https://example.com/api/data"))
time.sleep(12)
print(fetch_data_from_api("https://example.com/api/data"))

输出 (大致):

Fetching data from https://example.com/api/data...
Function executed and value cached. Last updated: 1678886400.0
Data from https://example.com/api/data at 1678886400.0
Value retrieved from cache.
Data from https://example.com/api/data at 1678886400.0
Fetching data from https://example.com/api/data...
Function executed and value cached. Last updated: 1678886417.0
Data from https://example.com/api/data at 1678886417.0

在这个例子中，cache 装饰器接收 max_age 参数，用于设置缓存的有效时间。如果缓存不存在或已过期，则执行被装饰的函数，并将结果缓存起来。 否则，直接从缓存中返回结果。 functools.wraps 确保了被装饰的函数 fetch_data_from_api 的元数据（例如 __name__ 和 __doc__）被正确保留。

2. 类装饰器

除了函数装饰器，我们还可以使用类来作为装饰器。类装饰器可以用来维护状态，并且可以提供更复杂的逻辑。

2.1 实现原理

类装饰器是通过实现 __call__ 方法来实现的。当使用类作为装饰器时，Python 会调用该类的构造函数，并将被装饰的函数作为参数传递给构造函数。然后，当调用被装饰的函数时，Python 会调用该类的 __call__ 方法。

2.2 示例代码

class CallCount:
    def __init__(self, func):
        self.func = func
        self.count = 0

    def __call__(self, *args, **kwargs):
        self.count += 1
        print(f"Call count: {self.count}")
        return self.func(*args, **kwargs)

@CallCount
def say_hello(name):
    print(f"Hello, {name}!")

say_hello("Bob")
say_hello("Charlie")
say_hello("David")

输出:

Call count: 1
Hello, Bob!
Call count: 2
Hello, Charlie!
Call count: 3
Hello, David!

2.3 代码解释

  • CallCount(func): 当使用 @CallCount 装饰 say_hello 函数时，Python 会创建一个 CallCount 类的实例，并将 say_hello 函数作为参数传递给构造函数 __init__。
  • __call__(*args, **kwargs): 当调用 say_hello("Bob") 时，Python 会调用 CallCount 实例的 __call__ 方法。该方法会增加计数器 count，并调用原始函数 func。

2.4 带参数的类装饰器

和函数装饰器一样，类装饰器也可以带参数。 这需要我们在类的 __init__ 方法中接收参数，并将被装饰的函数保存在类的属性中。

class Trace:
    def __init__(self, log_file):
        self.log_file = log_file

    def __call__(self, func):
        def wrapper(*args, **kwargs):
            with open(self.log_file, "a") as f:
                f.write(f"Calling function {func.__name__} with args {args} and kwargs {kwargs}n")
            return func(*args, **kwargs)
        return wrapper

@Trace("trace.log")
def add(x, y):
    return x + y

add(2, 3)
add(5, 7)

在这个例子中，Trace 类接收 log_file 参数，用于指定日志文件的路径。 __call__ 方法接收被装饰的函数 func，并返回包装后的函数 wrapper。 wrapper 函数在调用原始函数之前，会将函数名、参数和关键字参数写入日志文件。

2.5 使用场景

类装饰器在以下场景中非常有用：

  • 状态管理: 需要维护状态的装饰器，例如计数器、缓存。
  • 复杂的逻辑: 需要执行复杂逻辑的装饰器，例如事务管理、AOP。
  • 可配置性: 需要根据配置来改变行为的装饰器，例如日志级别、权限控制。

2.6 另一个更复杂例子 – 限制函数调用频率

import time

class RateLimit:
    def __init__(self, calls_per_second):
        self.calls_per_second = calls_per_second
        self.last_called = 0.0

    def __call__(self, func):
        def wrapper(*args, **kwargs):
            now = time.time()
            time_since_last_call = now - self.last_called

            if time_since_last_call < (1.0 / self.calls_per_second):
                time.sleep((1.0 / self.calls_per_second) - time_since_last_call)

            self.last_called = time.time()
            return func(*args, **kwargs)
        return wrapper

@RateLimit(calls_per_second=2) # 限制每秒最多调用两次
def expensive_operation(data):
    print(f"Processing data: {data}")
    time.sleep(0.4) # 模拟耗时操作

for i in range(5):
    expensive_operation(i)

在这个例子中，RateLimit 类接收 calls_per_second 参数，用于设置每秒最多调用的次数。wrapper 函数会检查上次调用到现在的时间间隔，如果时间间隔小于 1.0 / self.calls_per_second，则会暂停一段时间，以确保函数调用频率不超过限制。

3. wraps 的正确使用

在使用装饰器时，一个常见的问题是被装饰的函数的元数据（例如 __name__、__doc__）会丢失，被替换成包装函数的元数据。 这会导致调试和文档生成出现问题。 functools.wraps 可以解决这个问题。

3.1 作用

functools.wraps 是一个装饰器，它可以将被装饰的函数的元数据复制到包装函数。

3.2 示例代码

import functools

def my_decorator(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        """Wrapper function docstring"""
        print("Before the function call.")
        result = func(*args, **kwargs)
        print("After the function call.")
        return result
    return wrapper

@my_decorator
def say_hello(name):
    """Say hello to someone."""
    print(f"Hello, {name}!")

print(say_hello.__name__)
print(say_hello.__doc__)

输出:

say_hello
Say hello to someone.

如果没有使用 functools.wraps，输出将会是：

wrapper
Wrapper function docstring

3.3 代码解释

  • @functools.wraps(func): 这个装饰器会将 func 的 __name__、__doc__ 等元数据复制到 wrapper 函数。

3.4 最佳实践

在使用装饰器时，务必使用 functools.wraps 来保留原始函数的元数据。这不仅可以提高代码的可读性和可维护性，还可以避免一些潜在的问题。

3.5 结合之前的缓存装饰器

之前缓存装饰器例子已经使用了functools.wraps，是为了保证被装饰函数fetch_data_from_api的元信息能正确传递。

4. 总结，提升装饰器的理解水平

我们深入探讨了 Python 装饰器的高级用法，包括带参数的装饰器、类装饰器以及 wraps 的正确使用。掌握这些技巧可以让你编写更灵活、更强大、更易于维护的装饰器。 记住使用 functools.wraps 来保留原始函数的元数据，避免一些潜在的问题。

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注