如何向Python装饰器传递额外的参数？

Question

我有一个如下的装饰器。

def myDecorator(test_func):
    return callSomeWrapper(test_func)
def callSomeWrapper(test_func):
    return test_func
@myDecorator
def someFunc():
    print 'hello'

我想要增强这个装饰器，使其能够接受另一个参数，就像下面这样

def myDecorator(test_func,logIt):
    if logIt:
        print "Calling Function: " + test_func.__name__
    return callSomeWrapper(test_func)
@myDecorator(False)
def someFunc():
    print 'Hello'

但是这段代码会报错：

TypeError: myDecorator()需要传入两个参数，但只传了1个参数

为什么函数没有自动传递？如何显式地将函数传递给装饰器函数？

- balki

---

@KitHo - 这是一个布尔标志，因此使用布尔值是正确的方法。 - AKX

@gd是什么意思？是“good”吗？ - Rob Bednark

这回答您的问题吗？带参数的装饰器？ - Abdul Aziz Barkat

Answer 1

由于您像调用函数一样调用装饰器，所以它需要返回另一个函数，这个函数才是实际的装饰器：

def my_decorator(param):
    def actual_decorator(func):
        print("Decorating function {}, with parameter {}".format(func.__name__, param))
        return function_wrapper(func)  # assume we defined a wrapper somewhere
    return actual_decorator

外部函数将接收您明确传递的任何参数，并应返回内部函数。内部函数将接收要装饰的函数，并返回修改后的函数。

通常，您希望装饰器通过将其包装在一个包装函数中来更改函数行为。以下是一个示例，可根据需要在调用函数时添加日志记录：

def log_decorator(log_enabled):
    def actual_decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            if log_enabled:
                print("Calling Function: " + func.__name__)
            return func(*args, **kwargs)
        return wrapper
    return actual_decorator

functools.wraps函数会将原函数的名称和文档字符串等内容复制到包装函数中，使其更类似于原函数。

示例用法:

>>> @log_decorator(True)
... def f(x):
...     return x+1
...
>>> f(4)
Calling Function: f
5

- interjay

---

使用functools.wraps是明智之举——它保留了被装饰函数的原始名称、文档字符串等内容。 - AKX

@AKX：谢谢，我已将此添加到第二个示例中。 - interjay

所以基本上装饰器总是只接受一个参数，即函数。但是装饰器可以是一个函数的返回值，该函数可能会带有参数。这正确吗？ - balki

没错，那是正确的。让事情变得混乱的是许多人也会将外部函数（在这里是 myDecorator）称为装饰器。这对于装饰器的用户来说很方便，但在你尝试编写一个装饰器时可能会令人困惑。 - interjay

令我困惑的小细节：如果您的 log_decorator 函数带有默认参数，那么您不能使用 @log_decorator 这种方式进行装饰，而必须使用 @log_decorator()。 - Stardidi

Answer 2

只是提供一个不同的观点：

这个语法

@expr
def func(...): #stuff

等价于

def func(...): #stuff
func = expr(func)

特别地，expr 可以是任何你喜欢的东西，只要它能够被调用。 特别地，expr 可以是一个装饰器工厂：你给它一些参数，它就给你一个装饰器。因此，也许更好地理解您的情况是

dec = decorator_factory(*args)
@dec
def func(...):

这可以被缩短为

@decorator_factory(*args)
def func(...):

当然，由于`decorator_factory`看起来像一个装饰器，人们往往会以反映这一点的名称来命名它。当你尝试跟踪间接级别时，这可能会令人困惑。

- Katriel

---

谢谢，这确实帮助我理解背后的原理。 - Aditya Sriram

Answer 3

我只想添加一些有用的技巧，可以使装饰器参数变为可选。这还可以重复使用装饰器并减少嵌套。

import functools

def myDecorator(test_func=None,logIt=None):
    if test_func is None:
        return functools.partial(myDecorator, logIt=logIt)
    @functools.wraps(test_func)
    def f(*args, **kwargs):
        if logIt==1:
            print 'Logging level 1 for {}'.format(test_func.__name__)
        if logIt==2:
            print 'Logging level 2 for {}'.format(test_func.__name__)
        return test_func(*args, **kwargs)
    return f

#new decorator 
myDecorator_2 = myDecorator(logIt=2)

@myDecorator(logIt=2)
def pow2(i):
    return i**2

@myDecorator
def pow3(i):
    return i**3

@myDecorator_2
def pow4(i):
    return i**4

print pow2(2)
print pow3(2)
print pow4(2)

- Alexey Smirnov

---

抱歉，我又顶起了这个帖子。请问这种模式被认为是一个好的实践吗？ - tahesse

Answer 4

只是另一种装饰器的实现方式。我发现这种方式最容易理解。

class NiceDecorator:
    def __init__(self, param_foo='a', param_bar='b'):
        self.param_foo = param_foo
        self.param_bar = param_bar

    def __call__(self, func):
        def my_logic(*args, **kwargs):
            # whatever logic your decorator is supposed to implement goes in here
            print('pre action baz')
            print(self.param_bar)
            # including the call to the decorated function (if you want to do that)
            result = func(*args, **kwargs)
            print('post action beep')
            return result

        return my_logic

# usage example from here on
@NiceDecorator(param_bar='baaar')
def example():
    print('example yay')


example()

- Robert Fey

---

谢谢！我看了大约30分钟的一些令人费解的“解决方案”，但这是第一个真正有意义的。 - canhazbits

非常聪明的方法，可以实现具有可选参数的装饰器，而无需创建嵌套复杂性。将我的所有自定义装饰器重构为此，谢谢！ - Caio Castro

唯一的缺点是，当你只想使用默认参数时，它不能在没有显式调用装饰器的情况下使用。即使你不传递任何参数，你也必须像这样装饰：@NiceDecorator()。我想知道是否可以修改它，以便在不传递参数时可以写成 @NiceDecorator？ - UpTheIrons

唯一的缺点是，当您只想使用默认参数时，它不能在没有显式调用装饰器的情况下使用。即使您不传递任何参数，您也必须像这样进行装饰 @NiceDecorator()。我想知道是否可以修改它以使传递零个参数时编写 @NiceDecorator 成为可能？ - UpTheIrons

Answer 5

现在，如果您想使用装饰器decorator_with_arg调用一个名为function1的函数，并且在这种情况下，函数和装饰器都带有参数，

def function1(a, b):
    print (a, b)

decorator_with_arg(10)(function1)(1, 2)

- SuperNova

Answer 6

接受多个输入的装饰器就像{{link1：dataclasses装饰器}}。

在这个例子中，dataclass接受三种语法：

@dataclass
class C:
    ...

@dataclass()
class C:
    ...

@dataclass(init=True, repr=True, eq=True, order=False, unsafe_hash=False, frozen=False,
           match_args=True, kw_only=False, slots=False, weakref_slot=False)
class C:
    ...

在创建具有相同行为的装饰器时，您可以使用以下代码：

import inspect

def customdecorator(*args, **kwargs):
    def decorator(func):
        print('input decorator:', args, kwargs)
        def __wrapper(*func_args, **func_kwargs):
            print('input decorator inside function:', args, kwargs)
            print('input function:', func_args, func_kwargs)
            # Do something before calling the function
            result = func(*func_args, **func_kwargs)
            # Do something after calling the function
            return result
        return __wrapper

    print('input root:', args, kwargs)
    if len(kwargs) > 0:
        # Decorator is used with arguments, e.g., @functionmethod(arg1=val1, arg2=val2)
        return decorator
    if len(args) == 0:
        return decorator
    if len(args) == 1:
        return decorator(args[0])

# Example usages
@customdecorator
def example1():
    print("Function without call")

@customdecorator()
def example2():
    print("Function without arguments")

@customdecorator(arg1="value1", arg2="value2")
def example3(arg2):
    print(f"Function with arguments: {arg2}")

example1()
example2()
example3(arg2="ex2")

在你的情况下，这将是

def myDecorator(*args, **kwargs):
    def decorator(func):
        def __wrapper(*func_args, **func_kwargs):
            # Do something before calling the function
            test_func = kwargs.get('test_func', None)
            logIt = kwargs.get('logIt', None)
            if logIt:
                print("Calling Function: " + test_func.__name__)

            result = func(*func_args, **func_kwargs)

            # Do something after calling the function
            return result
        return __wrapper

    if len(kwargs) > 0:
        # Decorator is used with arguments, e.g., @functionmethod(arg1=val1, arg2=val2)
        return decorator
    if len(args) == 0:
        return decorator
    if len(args) == 1:
        return decorator(args[0])


@myDecorator(logIt=False)
def someFunc():
    print('Hello')

someFunc()

注意事项是：

1. 可选参数必须使用关键字参数。
2. 默认值通过 kwargs.get(..., ...) 输入。

- Muhammad Yasirroni