python标准库functools 装饰器

1. 装饰器的定义

装饰器是一种函数，它能够动态地修改或增强其他函数或方法的行为，而无需直接修改它们的代码。functools 中的装饰器可以帮助简化代码逻辑、提高代码复用性和效率。

本质上是一个接受函数作为输入并返回另一个函数的函数。它允许你在不改变原函数的前提下，为函数添加功能。例如，可以为某个函数添加缓存机制、调试信息、或执行时间记录等功能。

2. functools 提供的装饰器

2.1 @lru_cache

@lru_cache：用于缓存函数的返回结果，避免重复计算。通过缓存函数结果来提高函数效率，适合处理重复计算的场景。nnUnet 计算guassian 重要性图时使用到。

from functools import lru_cache

@lru_cache(maxsize=2)
def compute_something(x):
    # 模拟占据内存的昂贵计算
    print(f"Computing for {x}...")
    return x * x

# 调用
print(compute_something(2))  # 第一次计算，输出: Computing for 2... \n 4
print('*' *50)
print(compute_something(2))  # 第二次调用，直接从缓存中获取，输出: 4
print('*' *50)
print(compute_something(3))  # 第一次计算，输出: Computing for 3... \n 9
print('*' *50)
print(compute_something(2))  # 缓存中仍有结果，输出: 4
print('*' *50)
print(compute_something(4))  # 缓存已满，计算新的，输出: Computing for 4... \n 16
print('*' *50)
print(compute_something(3))  # 缓存被替换，重新计算，输出: Computing for 3... \n 9
print('*' *50)

2.2 @wraps

@wraps：用于编写装饰器时保持被装饰函数的原始属性。通常，编写装饰器时，原函数的元数据（如函数名、注释文档等）可能会被替换为装饰器函数的信息。@wraps 可以解决这个问题，确保原函数的元数据不会丢失。

from functools import wraps

def my_decorator(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        print("Before the function is called")
        result = func(*args, **kwargs)
        print("After the function is called")
        return result
    return wrapper

@my_decorator
def my_function():
    """This is my function."""
    print("Executing the function")

my_function()
print(my_function.__name__)  # 输出 'my_function'
print(my_function.__doc__)   # 输出 'This is my function.'

2.3 @total_ordering

@total_ordering：为类实现完整的比较运算符。只需定义一个或两个基础比较方法（如 __lt__ 和 __eq__），其他比较运算符（如 __gt__, __le__, __ge__, __ne__）将自动生成。

from functools import total_ordering

@total_ordering
class Person:
    def __init__(self, name, age):
        self.name = name
        self.age = age

    def __eq__(self, other):
        return self.age == other.age

    def __lt__(self, other):
        return self.age < other.age

p1 = Person("Alice", 30)
p2 = Person("Bob", 25)
print(p1 > p2)  # True，自动生成 __gt__

2.4 @partial

partial 并不直接是装饰器，但作用类似，允许部分应用一个函数的参数，创建一个新的函数。这个新的函数会有一些预设的参数值，方便复用。

from functools import partial

def power(base, exponent):
    return base ** exponent

square = partial(power, exponent=2)
cube = partial(power, exponent=3)

print(square(4))  # 16
print(cube(2))    # 8