如何创建既可以带参数也可以不带参数的装饰器?

2022-04-07 00:00:00 python decorator

问题描述

我想创建一个Python修饰器,它可以与参数一起使用:

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

或不使用它们(例如,默认情况下将输出重定向到stderr):

@redirect_output
def foo():
    ....

这完全可能吗?

请注意,我不是在寻找重定向输出问题的不同解决方案,这只是我希望实现的语法的一个示例。


解决方案

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

正如Thobe的答案所说,处理这两种情况的唯一方法是检查两种情况。最简单的方法是简单地检查是否有一个参数并且它是Callabe(注意:如果您的修饰器只有一个参数,并且它恰好是一个可调用的对象,则需要进行额外的检查):

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)

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

在第二种情况下,您返回一个‘new’修饰符,它以某种方式使用了通过*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

相关文章