博客

Python 向装饰器添加参数

  • Post category:Python

在 Python 中，装饰器是一种用于修改函数或类行为的语法结构。装饰器本质上是一个函数，它可以接收一个函数作为参数，对该函数进行包装后返回一个新的函数。在装饰器中可以对传入的函数进行增强功能、修饰输出结果、添加前置/后置处理逻辑等。在某些场合下，我们需要向装饰器传递参数以便定制化装饰器的行为，下面将详细讲解Python向装饰器添加参数的使用方法。

装饰器的基本使用

在讲解如何向装饰器添加参数之前，有必要回顾一下 Python 的装饰器基本用法。下面给出一个装饰器的例子，该装饰器实现在函数执行前打印一条日志信息：

def logger(func):
    def wrapper(*args, **kwargs):
        print('Function "{}" is called'.format(func.__name__))
        return func(*args, **kwargs)
    return wrapper

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

print(add(1, 2))

运行上面的代码，可以得到如下输出：

Function "add" is called
3

上面的代码中，logger 是一个装饰器，它接收一个函数 func 作为参数，返回一个被包装后的函数 wrapper。在 wrapper 函数内部，先输出一条日志信息，然后再调用原来的 func 函数，并返回其返回值。通过在 add 函数定义前加上 @logger，实现了在执行 add 函数前输出日志的功能。

向装饰器添加参数的方法

在某些场合下，我们需要向装饰器传递一些参数以便定制其行为。可以通过在定义装饰器时增加一个函数，在这个函数中传入装饰器的参数，最后返回原装饰器的方式来实现。

def my_decorator(color):
    def decorator(func):
        def wrapper(*args, **kwargs):
            print('{} color is used.'.format(color))
            return func(*args, **kwargs)
        return wrapper
    return decorator

@my_decorator(color='red')
def print_hello(name):
    print('Hello, {}!'.format(name))

print_hello('Peter')

上面的代码实现了一个颜色装饰器，装饰器用来设置输出文字使用的颜色。注意，my_decorator 这个函数实际上是一个工厂函数，它返回了一个真正的装饰器 decorator。decorator 也是一个函数，它接收一个函数 func 作为参数，返回一个被包装后的函数 wrapper。在 wrapper 函数中，输出使用的颜色，然后再调用原来的 func 函数，并返回其返回值。通过在 print_hello 函数定义前加上 @my_decorator(color='red')，实现了用红色输出 Hello, XXX! 的功能。

除了以上的方法，我们还可以通过使用 functools 模块中的 wraps 函数来实现。wraps 函数可以保留原函数的元信息，如函数名等，避免在使用装饰器包装完函数之后丢失属性信息的问题。

from functools import wraps

def my_decorator(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        print('Before')
        return func(*args, **kwargs)
    return wrapper

@my_decorator
def hello_world():
    """A simple hello world function"""
    print('Hello, World!')

print(hello_world.__name__, hello_world.__doc__)

输出：

hello_world A simple hello world function

上面的代码实现了一个简单的装饰器，输出 Before，然后再执行原来的 hello_world 函数。通过 wraps 函数，我们保留了 hello_world 函数原本的名称和文档字符串信息。

Tags: