如何创建可以使用参数或不使用参数的Python装饰器？

Question

我想创建一个Python decorator，它可以与参数一起使用：

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

或者不使用它们（例如，默认情况下将输出重定向到stderr）：

@redirect_output
def foo():
    ....

那有可能吗？

注意，对于重定向输出的问题，我并没有寻找不同的解决方案，这只是我想要实现的语法的一个例子。

Answer 1

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

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")语法时，redirect_output用一个参数"output.log"调用，它必须返回一个decorator，该decorator接受要作为参数修饰的函数。当用作@redirect_output时，直接调用它，并将函数修饰为参数。

或者换句话说，@语法后面必须跟一个表达式，其结果是一个函数接受一个要修饰的函数作为其唯一参数，并返回修饰的函数。表达式本身可以是函数调用，这是@redirect_output("output.log")的情况。令人费解，但却是事实：-）

Answer 2

我知道这个问题已经过时了，但有些评论是新的，虽然所有可行的解决方案基本上都是一样的，但大多数都不是很干净，也不容易阅读。

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

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)

在第一种情况下，您可以像任何普通的decorator那样，返回传入函数的修改或包装版本。

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

这是很好的，所有，但必须写出来的每一个装饰你可以很烦人，而不是干净。相反，它将是很好的，能够自动修改我们的装饰，而不必重新编写他们。。。但这就是装饰师的职责！

使用下面的decorator decorator，我们可以对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来装饰我们的装饰人员，他们将在有无争论的情况下工作，但有一个警告：

我在上面提到过，但这里应该重复一次，这个decorator中的check对decorator可以接收的参数做了一个假设（即它不能接收一个可调用的参数）。由于我们现在正在使它适用于任何发电机，它需要记住，或修改，如果它将矛盾。

以下说明其用途：

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

Answer 3

使用具有默认值的关键字参数（如kquinn所建议的）是一个好主意，但需要包含括号：

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

如果您想要一个在装饰器上没有括号的版本，那么您必须在装饰器代码中考虑这两种情况。

如果使用的是Python3.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)

在Python2.x中，这可以用varargs技巧模拟：

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():
    ...