如何创建一个可以带参数或不带参数的装饰器？

Question

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

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

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

 @redirect_output
def foo():
    ....

这可能吗？

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

原文由 elifiner 发布，翻译遵循 CC BY-SA 4.0 许可协议

Answer 1

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

就像 thobe 的回答所说，处理这两种情况的唯一方法是检查这两种情况。最简单的方法是简单地检查是否有一个参数并且它是 callabe（注意：如果你的装饰器只接受 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 传入的信息。

这很好，但是必须为你制作的每个装饰器写出来可能会很烦人而且不那么干净。相反，能够自动修改我们的装饰器而不必重写它们会很好……但这就是装饰器的用途！

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

 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 发布，翻译遵循 CC BY-SA 3.0 许可协议

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 发布，翻译遵循 CC BY-SA 2.5 许可协议