如何创建一个装饰器，使其可以带或不带参数使用？

Question

我想创建一个Python装饰器，可以使用参数：

@redirect_output("somewhere.log")
def foo():
    ....

或者没有它们（例如默认将输出重定向到stderr）：

@redirect_output
def foo():
    ....

这有可能吗？

请注意，我不是在寻找重定向输出问题的其他解决方案，只是想实现类似语法的示例。

- gooli

---

默认的@redirect_output看起来非常不具有信息性。我建议这样做是一个糟糕的想法。使用第一种形式会大大简化你的生活。 - S.Lott

虽然这是一个有趣的问题，但在我看到它并查阅了文档之前，我本以为 @f 和 @f() 是一样的，而且我仍然认为它应该是一样的（任何提供的参数都将被附加到函数参数上）。 - rog

这个装饰器工厂/装饰器模式很好，第一个默认参数function=None，我会进一步将其余参数设为关键字参数。 - Tomasz Gandor

Answer 1

我知道这个问题很旧，但一些评论是新的，虽然所有可行的解决方案本质上都相同，但大多数解决方案不太干净或易于阅读。

就像thobe的答案所说，处理这两种情况的唯一方法是检查这两种情况。 最简单的方法就是简单地检查是否有单个参数且它是可调用的（注意：如果您的装饰器只接受1个参数并且恰好是可调用对象，则需要进行额外的检查）：

def decorator(*args, **kwargs):
    if len(args) == 1 and len(kwargs) == 0 and callable(args[0]):
        # called as @decorator
    else:
        # called as @decorator(*args, **kwargs)

在第一种情况下，你会像任何普通装饰器一样返回经过修改或包装的传入函数。

在第二种情况下，你会返回一个“新”的装饰器，它以某种方式使用了通过 *args、**kwargs 传入的信息。

这是完全没问题的，但是每次创建装饰器都必须写出这些内容可能会相当烦人，并且不够简洁。相反，能够自动修改我们的装饰器而无需重新编写它们会很好，但这就是装饰器的作用！

使用以下装饰器装饰器，我们可以装饰我们的装饰器，使它们可以带或不带参数使用：

def doublewrap(f):
    '''
    a decorator decorator, allowing the decorator to be used as:
    @decorator(with, arguments, and=kwargs)
    or
    @decorator
    '''
    @wraps(f)
    def new_dec(*args, **kwargs):
        if len(args) == 1 and len(kwargs) == 0 and callable(args[0]):
            # actual decorated function
            return f(args[0])
        else:
            # decorator arguments
            return lambda realf: f(realf, *args, **kwargs)

    return new_dec

现在，我们可以使用 @doublewrap 装饰我们的装饰器，它们将适用于带参数和不带参数的情况，但有一个注意点：

我之前已经提到过，但还是要重复说明一下，这个装饰器中的检查做出了一个假设，即一个装饰器不能接收单个可调用参数。由于我们现在要将其应用于任何生成器，因此需要记住它，或者在将来修改，以防止相冲突。

以下演示了它的用法：

def test_doublewrap():
    from util import doublewrap
    from functools import wraps    

    @doublewrap
    def mult(f, factor=2):
        '''multiply a function's return value'''
        @wraps(f)
        def wrap(*args, **kwargs):
            return factor*f(*args,**kwargs)
        return wrap

    # try normal
    @mult
    def f(x, y):
        return x + y

    # try args
    @mult(3)
    def f2(x, y):
        return x*y

    # try kwargs
    @mult(factor=5)
    def f3(x, y):
        return x - y

    assert f(2,3) == 10
    assert f2(2,5) == 30
    assert f3(8,1) == 5*7

- bj0

---

如果第一个参数是类，callable(args[0]) 将无法起作用。 你可以使用 isinstance(args[0], types.FunctionType) 来检测类。 - Noam Nol

我提到这是我提供的示例中的假设之一。对于像传递类这样的特殊情况，您需要修改它以适应您的情况。 - bj0

Answer 2

使用具有默认值的关键字参数（正如kquinn所建议的那样）是个好主意，但需要在调用时加上括号：

@redirect_output()
def foo():
    ...

如果您希望获得一个不需要装饰器中的括号也能正常运行的版本，您需要在装饰器代码中考虑这两种情况。

如果您使用的是Python 3.0版本，您可以使用关键字参数来实现此功能:

def redirect_output(fn=None,*,destination=None):
  destination = sys.stderr if destination is None else destination
  def wrapper(*args, **kwargs):
    ... # your code here
  if fn is None:
    def decorator(fn):
      return functools.update_wrapper(wrapper, fn)
    return decorator
  else:
    return functools.update_wrapper(wrapper, fn)

在Python 2.x中，可以使用可变参数技巧来模拟此操作：

def redirected_output(*fn,**options):
  destination = options.pop('destination', sys.stderr)
  if options:
    raise TypeError("unsupported keyword arguments: %s" % 
                    ",".join(options.keys()))
  def wrapper(*args, **kwargs):
    ... # your code here
  if fn:
    return functools.update_wrapper(wrapper, fn[0])
  else:
    def decorator(fn):
      return functools.update_wrapper(wrapper, fn)
    return decorator

任何一个版本都允许您编写以下代码：

@redirected_output
def foo():
    ...

@redirected_output(destination="somewhere.log")
def bar():
    ...

- thobe

---

你会在 your code here 中放什么？你如何调用被装饰的函数？fn(*args, **kwargs) 无法正常工作。 - lum

我认为有一个更简单的答案，创建一个类作为带选项参数的装饰器。再创建一个具有相同默认参数的函数并返回装饰器类的新实例。 应该类似于：def f(a=5): return MyDecorator(a=a)和class MyDecorator(object): def __init__(self, a=5): ....很抱歉在评论中写出来有些困难，但我希望这足够简单易懂。 - Omer Ben Haim

Answer 3

我知道这是一个老问题，但我真的不喜欢任何提出的技术，所以我想添加另一种方法。我注意到django在他们的django.contrib.auth.decorators中的login_required装饰器中使用了一种非常简洁的方法。就像你可以在装饰器的文档中看到的那样，它可以单独使用 @login_required 或带有参数的形式，例如@login_required(redirect_field_name='my_redirect_field')。

他们的实现方式非常简单。在装饰器的参数之前，他们加入了一个kwarg参数(function=None)。如果装饰器单独使用，function将是它所装饰的实际函数；而如果它带有参数，则function将是None。

示例：

from functools import wraps

def custom_decorator(function=None, some_arg=None, some_other_arg=None):
    def actual_decorator(f):
        @wraps(f)
        def wrapper(*args, **kwargs):
            # Do stuff with args here...
            if some_arg:
                print(some_arg)
            if some_other_arg:
                print(some_other_arg)
            return f(*args, **kwargs)
        return wrapper
    if function:
        return actual_decorator(function)
    return actual_decorator

---

@custom_decorator
def test1():
    print('test1')

>>> test1()
test1

---

@custom_decorator(some_arg='hello')
def test2():
    print('test2')

>>> test2()
hello
test2

---

@custom_decorator(some_arg='hello', some_other_arg='world')
def test3():
    print('test3')

>>> test3()
hello
world
test3

我发现Django使用的这种方法比这里提出的任何其他技术更加优雅且易于理解。

- dgel

---

是的，我喜欢这个方法。请注意，在调用装饰器时，您必须使用kwargs，否则第一个位置参数将分配给“function”，然后事情会破裂，因为装饰器尝试调用该第一个位置参数，就好像它是您装饰的函数。 - Dustin Wyatt

是的，第一个参数不是 kwarg，它是一个带有默认值的定位参数。但是您可以使其余的参数仅限于关键字参数。 - Tomasz Gandor

Answer 4

已经有几个回答很好地解决了你的问题。不过，关于样式，我更喜欢使用 functools.partial 来解决这个装饰器问题，正如 David Beazley 的 Python Cookbook 3 所建议的那样：

from functools import partial, wraps

def decorator(func=None, foo='spam'):
    if func is None:
         return partial(decorator, foo=foo)

    @wraps(func)
    def wrapper(*args, **kwargs):
        # do something with `func` and `foo`, if you're so inclined
        pass
    
    return wrapper

虽然是的，你可以这样做

@decorator()
def f(*args, **kwargs):
    pass

没有花哨的变通方法，我觉得这看起来很奇怪，因此我喜欢有简单装饰选项@decorator。

至于第二个任务目标，将函数的输出重定向在这个Stack Overflow帖子中有所涉及。

---

如果你想深入了解，可以查看《Python Cookbook 3》的第9章（元编程），该书可以免费在线阅读。

其中一些材料被演示了（还有更多！）在 Beazley 精彩的YouTube视频Python 3 Metaprogramming中。

- henrywallace

---

我最喜欢这种风格，谢谢分享。对于其他人，我要注意，如果你试图在“wrapper”内更改“foo”，可能会出现“UnboundLocalError”的错误，在这种情况下，你应该声明“nonlocal foo”（或者选择一个不同的本地变量名，如“bar”，并设置为“bar = foo”）。 参见：https://stackoverflow.com/a/57184656/1588795 - n8henrie

请注意，此解决方案假定调用代码始终使用关键字参数。在这种情况下，@decorator('foo')将无法按预期工作。 - Björn Pollex

@BjörnPollex 如果你想支持位置参数，只需将 if func is None: 更改为 if not callable(func):。 - undefined

Answer 5

你需要检测两种情况，例如使用第一个参数的类型，并相应地返回包装器（在没有参数时使用）或装饰器（在使用参数时使用）。

from functools import wraps
import inspect

def redirect_output(fn_or_output):
    def decorator(fn):
        @wraps(fn)
        def wrapper(*args, **args):
            # Redirect output
            try:
                return fn(*args, **args)
            finally:
                # Restore output
        return wrapper

    if inspect.isfunction(fn_or_output):
        # Called with no parameter
        return decorator(fn_or_output)
    else:
        # Called with a parameter
        return decorator

使用@redirect_output("output.log")语法时，必须以单个参数"output.log"调用redirect_output，并返回一个接受要装饰的函数作为参数的装饰器。如果用@redirect_output，则直接将要装饰的函数作为参数调用。

换句话说：@语法后面必须跟一个表达式，其结果是接受要装饰的函数作为唯一参数并返回已装饰函数的函数。表达式本身可以是一个函数调用，这就是@redirect_output("output.log")的情况。有点绕，但没错 :-)

- Remy Blank

Answer 6

Python装饰器的调用方式取决于是否传递参数。实际上，装饰器只是一个（语法上受限制的）表达式。

在第一个示例中：

@redirect_output("somewhere.log")
def foo():
    ....

函数redirect_output被调用时，给定的参数应该返回一个装饰器函数，这个装饰器函数本身又以foo作为参数被调用，最终期望返回最终装饰后的函数。

等效的代码如下：

def foo():
    ....
d = redirect_output("somewhere.log")
foo = d(foo)

你的第二个示例的等效代码如下所示：

def foo():
    ....
d = redirect_output
foo = d(foo)

因此，您可以按照自己的意愿进行操作，但并非完全无缝：

import types
def redirect_output(arg):
    def decorator(file, f):
        def df(*args, **kwargs):
            print 'redirecting to ', file
            return f(*args, **kwargs)
        return df
    if type(arg) is types.FunctionType:
        return decorator(sys.stderr, arg)
    return lambda f: decorator(arg, f)

除非您希望在装饰器的参数中使用函数，否则这应该是可以接受的。如果在参数中使用函数，则装饰器将错误地假定它没有参数。如果将此装饰用于不返回函数类型的另一个装饰上，则也会失败。

另一种方法是始终要求调用装饰器函数，即使没有传递参数。在这种情况下，您的第二个示例将如下所示：

@redirect_output()
def foo():
    ....

装饰器函数的代码应该像这样：

def redirect_output(file = sys.stderr):
    def decorator(file, f):
        def df(*args, **kwargs):
            print 'redirecting to ', file
            return f(*args, **kwargs)
        return df
    return lambda f: decorator(file, f)

- rog

Answer 7

既然没有人提到这一点，还有一种解决方案是利用可调用类，我觉得这更加优雅，特别是在装饰器比较复杂且希望将其拆分为多个方法（函数）的情况下。这个解决方案利用了__new__魔法方法来实现其他人指出的基本功能。首先检测装饰器的使用方式，然后相应地调整返回值。

class decorator_with_arguments(object):

    def __new__(cls, decorated_function=None, **kwargs):

        self = super().__new__(cls)
        self._init(**kwargs)

        if not decorated_function:
            return self
        else:
            return self.__call__(decorated_function)

    def _init(self, arg1="default", arg2="default", arg3="default"):
        self.arg1 = arg1
        self.arg2 = arg2
        self.arg3 = arg3

    def __call__(self, decorated_function):

        def wrapped_f(*args):
            print("Decorator arguments:", self.arg1, self.arg2, self.arg3)
            print("decorated_function arguments:", *args)
            decorated_function(*args)

        return wrapped_f

@decorator_with_arguments(arg1=5)
def sayHello(a1, a2, a3, a4):
    print('sayHello arguments:', a1, a2, a3, a4)

@decorator_with_arguments()
def sayHello(a1, a2, a3, a4):
    print('sayHello arguments:', a1, a2, a3, a4)

@decorator_with_arguments
def sayHello(a1, a2, a3, a4):
    print('sayHello arguments:', a1, a2, a3, a4)

如果装饰器带有参数，则等同于以下内容：

result = decorator_with_arguments(arg1=5)(sayHello)(a1, a2, a3, a4)

可以看到，参数arg1被正确传递给构造函数，并且装饰的函数被传递给__call__。

但是如果装饰器没有使用参数，则等同于：

result = decorator_with_arguments(sayHello)(a1, a2, a3, a4)

你看，在这种情况下，装饰函数直接传递给构造函数，并且完全省略了对__call__的调用。这就是为什么我们需要在__new__魔术方法中使用逻辑来处理这种情况的原因。

为什么我们不能使用__init__而不是__new__呢？原因很简单：Python禁止从__init__返回除None以外的任何其他值。 警告 这种方法有两个副作用。

- 它不会保留函数签名。 - 类型检查器会感到困惑。例如，mypy只允许从__new__返回相同类的实例，而不允许返回其他任何东西。请参见here。

- Majo

Answer 8

补充其他答案：

“是否有一种方法可以构建一个装饰器，既可以带参数使用也可以不带参数使用？”

没有通用的方法，因为目前Python语言中缺少检测这两种不同用法的功能。

然而，正如其他答案所指出的那样（例如bj0s），确实存在一个笨拙的解决方法，即检查接收到的第一个位置参数的类型和值（并检查是否没有其他参数具有非默认值）。如果您保证用户永远不会将可调用对象作为装饰器的第一个参数传递，则可以使用此解决方法。请注意，对于类装饰器来说也是相同的（在上一个句子中将可调用对象替换为类）。

为了确信以上内容，我进行了相当多的研究，甚至实现了一个名为decopatch的库，该库使用上述所有策略的组合（以及许多其他策略，包括内省）根据您的需求执行“最智能的解决方案”。它提供了两种模式：嵌套和平面。

在“嵌套模式”中，您总是返回一个函数。

from decopatch import function_decorator

@function_decorator
def add_tag(tag='hi!'):
    """
    Example decorator to add a 'tag' attribute to a function. 
    :param tag: the 'tag' value to set on the decorated function (default 'hi!).
    """
    def _apply_decorator(f):
        """
        This is the method that will be called when `@add_tag` is used on a 
        function `f`. It should return a replacement for `f`.
        """
        setattr(f, 'tag', tag)
        return f
    return _apply_decorator

当处于“flat mode”时，您的方法直接成为应用装饰器时将执行的代码。 它被注入了被装饰函数对象f：

from decopatch import function_decorator, DECORATED

@function_decorator
def add_tag(tag='hi!', f=DECORATED):
    """
    Example decorator to add a 'tag' attribute to a function.
    :param tag: the 'tag' value to set on the decorated function (default 'hi!).
    """
    setattr(f, 'tag', tag)
    return f

但说实话，最好的情况是不需要任何库，直接从Python语言中获得该功能。如果您像我一样认为今天Python语言不能提供对这个问题的简洁回答很可惜，请毫不犹豫地在Python错误追踪器中支持这个想法：https://bugs.python.org/issue36553！

- smarie

---

“than the above” 不是在stackoverflow上一个有用的短语。不同的人会随着时间的推移以不同的顺序查看答案。无法知道你所指的答案是什么。 - Bryan Oakley

感谢您发现了这个问题，@BryanOakley。确实，如果您觉得它有用并投票支持，那么上面的信息就会更少。我已经相应地编辑了该信息。 - smarie

Answer 9

实际上，@bj0的解决方案中的警告情况很容易检查：

def meta_wrap(decor):
    @functools.wraps(decor)
    def new_decor(*args, **kwargs):
        if len(args) == 1 and len(kwargs) == 0 and callable(args[0]):
            # this is the double-decorated f. 
            # Its first argument should not be a callable
            doubled_f = decor(args[0])
            @functools.wraps(doubled_f)
            def checked_doubled_f(*f_args, **f_kwargs):
                if callable(f_args[0]):
                    raise ValueError('meta_wrap failure: '
                                'first positional argument cannot be callable.')
                return doubled_f(*f_args, **f_kwargs)
            return checked_doubled_f 
        else:
            # decorator arguments
            return lambda real_f: decor(real_f, *args, **kwargs)

    return new_decor

这是一些针对此故障安全版本的 meta_wrap 的测试用例。

    @meta_wrap
    def baddecor(f, caller=lambda x: -1*x):
        @functools.wraps(f)
        def _f(*args, **kwargs):
            return caller(f(args[0]))
        return _f

    @baddecor  # used without arg: no problem
    def f_call1(x):
        return x + 1
    assert f_call1(5) == -6

    @baddecor(lambda x : 2*x) # bad case
    def f_call2(x):
        return x + 1
    f_call2(5)  # raises ValueError

    # explicit keyword: no problem
    @baddecor(caller=lambda x : 100*x)
    def f_call3(x):
        return x + 1
    assert f_call3(5) == 600

- Ainz Titor

Answer 10

这对我有用：

def redirect_output(func=None, /, *, output_log='./output.log'):
    def out_wrapper(func):
        def wrapper(*args, **kwargs):
            res = func(*args, **kwargs)
            print(f"{func.__name__} finished, output_log:{output_log}")
            return res

        return wrapper

    if func is None:
        return out_wrapper  # @redirect_output()
    return out_wrapper(func)  # @redirect_output


@redirect_output
def test1():
    print("running test 1")


@redirect_output(output_log="new.log")
def test2():
    print("running test 2")

test1()
print('-----')
test2()

- WangShengguang