如何为Python装饰器添加参数？基础装饰器传参疑问解析

问题背景 #

已经实现了基础Python装饰器：

def printArgs(function):
    def wrapper(*args, **kwargs):
        print(args, kwargs)
        return function(*args, **kwargs)
    return wrapper

class A:
    @printArgs
    def add(self, number1, number2):
        return number1 + number2

尝试创建可接收额外参数的装饰器时遇到错误：

def addXToResult(function, x):
    def wrapper(*args, **kwargs):
        return function(*args, **kwargs) + x
    return wrapper

使用@addXToResult(2)装饰类方法时，报错missing 1 required positional argument: 'x'；改用@addXToResult(A.add, 2)则出现NameError: name 'A' is not defined。

一、为什么基础装饰器会自动传入被装饰函数？ #

Python的@是装饰器语法糖，@decorator等价于被装饰函数 = decorator(被装饰函数)。

以基础装饰器为例：

@printArgs
def add(self, number1, number2):
    return number1 + number2

实际执行逻辑是：

add = printArgs(add)

Python会自动将下方的add函数作为参数传入printArgs，这就是基础装饰器能直接拿到被装饰函数的原因。

二、如何创建带额外参数的装饰器？ #

需要给装饰器嵌套一层外层函数，用来接收额外参数，再返回真正的装饰器函数。修改后的代码如下：

def addXToResult(x):
    # 外层函数接收额外参数x
    def decorator(function):
        # 内层decorator才是真正接收被装饰函数的装饰器
        def wrapper(*args, **kwargs):
            result = function(*args, **kwargs)
            return result + x
        return wrapper
    return decorator

使用示例 #

class A:
    @addXToResult(2)
    def add(self, number1, number2):
        return number1 + number2

a = A()
print(a.add(1, 2))  # 输出：5

逻辑说明 #

1. addXToResult(2)先执行，接收参数2后返回内层的decorator函数；
2. Python自动将A.add作为参数传入decorator；
3. decorator返回wrapper函数，最终A.add被替换为wrapper；
4. 调用a.add(1,2)时，实际执行wrapper，先调用原add得到结果3，再加2返回5。

原写法错误原因 #

• 使用@addXToResult(2)时，相当于先调用addXToResult(2)，此时只传了一个参数，而原函数需要function和x两个参数，因此报错缺少x；
• 使用@addXToResult(A.add, 2)时，类A还未完成定义（装饰器在类定义阶段执行），A.add尚未存在，因此触发NameError。

内容的提问来源于stack exchange，提问作者Xiidref