使用 functools.wraps 定义函数修饰器

在 Python 中，函数修饰器是一种非常强大的工具，可以用来增强或修改函数的行为。然而，在定义函数修饰器时，如果不正确地处理函数的元数据（如名称、文档字符串等），可能会导致一些意外的问题。functools.wraps 是一个非常有用的装饰器，它可以确保修饰器正确地保留被修饰函数的元数据。

本文将详细介绍如何使用 functools.wraps 来定义函数修饰器，并展示几个实际应用的例子。

1. 理解 functools.wraps

functools.wraps 是一个装饰器，它用于更新包装函数的元数据，使其与被包装的原始函数一致。具体来说，它会更新以下元数据：

函数名称 (__name__)
函数文档字符串 (__doc__)
函数注释 (__annotations__)
函数模块 (__module__)
函数属性字典 (__dict__)

2. 基本用法

首先来看一个简单的例子，不使用 functools.wraps 的情况。

python 复制代码

def simple_decorator(func):
    def wrapper(*args, **kwargs):
        print("Before function call")
        result = func(*args, **kwargs)
        print("After function call")
        return result
    return wrapper

@simple_decorator
def my_function():
    """This is a simple function."""
    print("Function called")

print(my_function.__name__)  # 输出: wrapper
print(my_function.__doc__)   # 输出: None

可以看到，my_function 的名称和文档字符串都被替换成了 wrapper 的元数据。接下来，我们使用 functools.wraps 来解决这个问题。

python 复制代码

from functools import wraps

def simple_decorator(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        print("Before function call")
        result = func(*args, **kwargs)
        print("After function call")
        return result
    return wrapper

@simple_decorator
def my_function():
    """This is a simple function."""
    print("Function called")

print(my_function.__name__)  # 输出: my_function
print(my_function.__doc__)   # 输出: This is a simple function.

现在，my_function 的名称和文档字符串都得到了正确的保留。

3. 实际应用

接下来，我们看几个实际的应用场景，展示如何使用 functools.wraps 来定义函数修饰器。

3.1 日志记录

python 复制代码

from functools import wraps

def log(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        print(f"Calling {func.__name__} with args: {args}, kwargs: {kwargs}")
        result = func(*args, **kwargs)
        print(f"Finished calling {func.__name__}")
        return result
    return wrapper

@log
def add(a, b):
    """Add two numbers and return the sum."""
    return a + b

result = add(5, 3)
print(result)          # 输出: 8
print(add.__name__)    # 输出: add
print(add.__doc__)     # 输出: Add two numbers and return the sum.

3.2 性能监控

python 复制代码

import time
from functools import wraps

def timeit(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        end_time = time.time()
        print(f"{func.__name__} took {end_time - start_time:.4f} seconds to run.")
        return result
    return wrapper

@timeit
def heavy_computation(n):
    """Perform a heavy computation."""
    total = 0
    for i in range(n):
        total += i
    return total

result = heavy_computation(1000000)
print(result)          # 输出: 499999500000
print(heavy_computation.__name__)  # 输出: heavy_computation
print(heavy_computation.__doc__)   # 输出: Perform a heavy computation.

3.3 输入验证

python 复制代码

from functools import wraps

def validate_input(func):
    @wraps(func)
    def wrapper(a, b):
        if not isinstance(a, int) or not isinstance(b, int):
            raise TypeError("Inputs must be integers")
        return func(a, b)
    return wrapper

@validate_input
def multiply(a, b):
    """Multiply two numbers."""
    return a * b

result = multiply(5, 3)
print(result)          # 输出: 15
print(multiply.__name__)  # 输出: multiply
print(multiply.__doc__)   # 输出: Multiply two numbers.

总结

functools.wraps 是一个非常有用的工具，可以确保函数修饰器正确地保留被修饰函数的元数据。通过使用 functools.wraps，我们可以避免因元数据丢失而导致的各种问题，使代码更加健壮和易读。

以上几个例子展示了如何使用 functools.wraps 来定义常见的函数修饰器，包括日志记录、性能监控和输入验证。希望这些示例能够帮助你在实际开发中更好地利用函数修饰器。