在 Python 开发中,经常会遇到需要对函数进行扩展或修改其行为的情况,但又不希望直接修改函数本身的代码。这时,包装函数(也称为装饰器)就派上了大用场。它是一种非常优雅且强大的工具,能够以一种简洁的方式增强函数的功能。

一、包装函数的基本概念

包装函数本质上是一个返回函数的高阶函数。它接收一个函数作为参数,对其进行包装,然后返回一个新的函数。通过使用 @ 符号,可以很方便地将包装函数应用到目标函数上。

在 Python 中,通常使用 functools.wraps 来保留被包装函数的元信息(如函数名、文档字符串等),以确保在使用装饰器后,函数的这些信息不会丢失。

二、应用场景:系统参数与用户参数的结合

代码示例中,包装函数被用来解决一个非常实用的问题:当系统中需要调用一个函数时,这个函数的参数既包括用户提供的参数,也包括系统配置的参数。如果直接在每次调用时手动添加这些参数,代码会变得冗长且难以维护。而通过使用包装函数,可以将系统参数封装起来,让用户在调用时只需传入自己的参数即可。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
from functools import wraps
import importlib

def force_think_status_false(func, _dict):
@wraps(func)
def wrapper(*args, **kwargs):
for key, value in _dict.items():
kwargs[key] = value
return func(*args, **kwargs)
return wrapper

# 示例使用
model_path = "some_module"
class_name = "ChatModel"
model = importlib.import_module(model_path)
chat = getattr(model, class_name)

func_dict = {
"role": "assistant",
"req_url": "https://api.example.com/chat",
"api_key": "your_api_key"
}
func_dict["think_model"] = False if model_config.get('think_model', False) == False else True

chat = force_think_status_false(chat, func_dict)

在这个例子中,force_think_status_false 是一个包装函数,它接收一个函数 func 和一个字典 _dict。在包装后的函数 wrapper 中,它将 _dict 中的键值对添加到 kwargs 中,然后调用原始的 func。这样,当用户调用 chat 时,系统参数会被自动添加进去,而用户只需要传入自己的参数即可。

三、其他包装函数的使用场景

(一)日志记录

在开发过程中,经常需要记录函数的调用情况,以便进行调试或监控。通过包装函数,可以轻松地为函数添加日志记录功能,而无需修改函数本身的代码。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
from functools import wraps
import logging

logging.basicConfig(level=logging.INFO)

def log_decorator(func):
@wraps(func)
def wrapper(*args, **kwargs):
logging.info(f"Calling function {func.__name__} with args: {args}, kwargs: {kwargs}")
result = func(*args, **kwargs)
logging.info(f"Function {func.__name__} returned {result}")
return result
return wrapper

@log_decorator
def add(a, b):
return a + b

print(add(3, 4))

在这个例子中,log_decorator 会记录函数的调用信息和返回值。当调用 add(3, 4) 时,日志会输出函数的名称、传入的参数以及返回的结果。

(二)性能测试

有时候,可能需要测试某个函数的执行时间,以评估其性能。包装函数可以方便地实现这一功能,而无需在函数内部添加计时代码。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
from functools import wraps
import time

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

@timing_decorator
def compute_sum(n):
return sum(range(n))

print(compute_sum(1000000))

在这个例子中,timing_decorator 会在函数执行前后记录时间,并计算出函数的执行时间。当调用 compute_sum(1000000) 时,它会输出函数的执行时间。

(三)权限验证

在开发 Web 应用或 API 时,通常需要对某些函数进行权限验证,以确保只有授权的用户才能调用它们。包装函数可以方便地实现这一功能。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
from functools import wraps

def auth_decorator(func):
@wraps(func)
def wrapper(*args, **kwargs):
user = get_current_user() # 假设有一个函数可以获取当前用户
if not user.is_authenticated:
raise PermissionError("You do not have permission to access this function.")
return func(*args, **kwargs)
return wrapper

@auth_decorator
def sensitive_data():
return "Sensitive information"

# 假设的用户类和获取当前用户函数
class User:
def __init__(self, is_authenticated):
self.is_authenticated = is_authenticated

def get_current_user():
return User(is_authenticated=True) # 示例用户已认证

print(sensitive_data())

在这个例子中,auth_decorator 会检查当前用户是否已认证。如果用户未认证,则抛出 PermissionError 异常。

(四)缓存结果

对于一些计算成本较高的函数,可以使用包装函数来缓存其结果,以提高性能。当函数被多次调用且参数相同时,可以直接返回缓存的结果,而无需重新计算。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
from functools import wraps

def cache_decorator(func):
cache = {}
@wraps(func)
def wrapper(*args):
if args in cache:
return cache[args]
result = func(*args)
cache[args] = result
return result
return wrapper

@cache_decorator
def fibonacci(n):
if n <= 1:
return n
return fibonacci(n - 1) + fibonacci(n - 2)

print(fibonacci(30))

在这个例子中,cache_decorator 使用一个字典 cache 来存储函数的参数和结果。当调用 fibonacci(30) 时,它会缓存计算结果,避免重复计算。

重点补充

与之后的案例有点不一样的地方在于,我使用了以一个方法叫做wraps,@wraps(func) 是 Python 中 functools 模块提供的一个非常有用的工具,它主要用于在定义装饰器时保留被装饰函数的元信息(如函数名、文档字符串等)。

一、@wraps(func) 的含义

在 Python 中,当使用装饰器对一个函数进行包装时,被装饰的函数的元信息(如函数名、文档字符串、参数列表等)会被覆盖。例如,如果不使用 @wraps(func),被装饰函数的 __name__ 属性会被改为装饰器函数的名称,文档字符串也会丢失。这可能会导致一些问题,比如在调试时无法正确识别函数,或者在某些依赖函数元信息的框架中出现问题。

@wraps(func) 的作用就是将被装饰函数的元信息复制到装饰器返回的函数中,从而保留这些信息。它是一个装饰器工厂,接收一个函数作为参数,并返回一个新的装饰器。

二、@wraps(func) 的使用方法

1. 不使用 @wraps(func) 的情况

首先,来看一个不使用 @wraps(func) 的例子:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
def my_decorator(func):
def wrapper(*args, **kwargs):
print("Something is happening before the function is called.")
result = func(*args, **kwargs)
print("Something is happening after the function is called.")
return result
return wrapper

@my_decorator
def say_hello(name):
"""Print a greeting message."""
print(f"Hello, {name}!")

say_hello("Alice")
print(say_hello.__name__)
print(say_hello.__doc__)

输出结果如下:

1
2
3
4
5
Something is happening before the function is called.
Hello, Alice!
Something is happening after the function is called.
wrapper
None

可以看到,say_hello.__name__ 变成了 wrappersay_hello.__doc__ 也变成了 None。这是因为装饰器返回的函数是 wrapper,而 wrapper 的元信息覆盖了 say_hello 的元信息。

2. 使用 @wraps(func) 的情况

接下来,使用 @wraps(func) 来解决这个问题:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
from functools import wraps

def my_decorator(func):
@wraps(func) # 使用 @wraps(func) 保留元信息
def wrapper(*args, **kwargs):
print("Something is happening before the function is called.")
result = func(*args, **kwargs)
print("Something is happening after the function is called.")
return result
return wrapper

@my_decorator
def say_hello(name):
"""Print a greeting message."""
print(f"Hello, {name}!")

say_hello("Alice")
print(say_hello.__name__)
print(say_hello.__doc__)

输出结果如下:

1
2
3
4
5
Something is happening before the function is called.
Hello, Alice!
Something is happening after the function is called.
say_hello
Print a greeting message.

可以看到,say_hello.__name__ 仍然是 say_hellosay_hello.__doc__ 也保留了原始的文档字符串。这是因为 @wraps(func)say_hello 的元信息复制到了 wrapper 函数中。

3. @wraps(func) 的具体作用

@wraps(func) 主要做了以下几件事:

  • 将被装饰函数的 __name__ 属性复制到装饰器返回的函数中。
  • 将被装饰函数的 __doc__ 属性(文档字符串)复制到装饰器返回的函数中。
  • 将被装饰函数的其他元信息(如参数列表等)也复制到装饰器返回的函数中。

这样,即使函数被装饰器包装,仍然可以通过 __name____doc__ 等属性获取到原始函数的信息。

三、为什么需要 @wraps(func)

  1. 调试方便:在调试代码时,能够正确识别函数名和文档字符串是非常重要的。如果函数名被覆盖为装饰器的名称,可能会导致混淆。
  2. 框架兼容性:某些框架或工具依赖函数的元信息来实现特定功能。例如,一些测试框架、文档生成工具等可能会根据函数的文档字符串或参数列表来生成文档或执行测试。如果不使用 @wraps(func),可能会导致这些工具无法正常工作。
  3. 代码可读性:保留函数的原始元信息可以提高代码的可读性,让其他开发者更容易理解代码的意图。