如何构建带可选参数的装饰器？

Question

我想做一个装饰器，这个装饰器可以带参数使用，也可以不带参数使用。大概是这样的：

class d(object):
    def __init__(self,msg='my default message'):
        self.msg = msg
    def __call__(self,fn):
        def newfn():
            print self.msg
            return fn()
        return newfn

@d('This is working')
def hello():
    print 'hello world !'

@d
def too_bad():
    print 'does not work'

在我的代码中，只有带参数的装饰器能正常工作。那么我该怎么做才能让这两种用法都能正常工作呢？

Answer 1

这样做是可以的。

def d(arg):
    if callable(arg):  # Assumes optional argument isn't.
        def newfn():
            print('my default message')
            return arg()
        return newfn
    else:
        def d2(fn):
            def newfn():
                print(arg)
                return fn()
            return newfn
        return d2

@d('This is working')
def hello():
    print('hello world !')

@d  # No explicit arguments will result in default message.
def hello2():
    print('hello2 world !')

@d('Applying it twice')
@d('Would also work')
def hello3():
    print('hello3 world !')

hello()
hello2()
hello3()

输出结果：

This is working
hello world !
my default message
hello2 world !
Applying it twice
Would also work
hello3 world !

如果一个装饰器函数用@调用时没有传入任何明确的参数，它会使用下面定义的函数来调用。如果传入了参数，那么它会先用这些参数进行调用，然后再用这个初步调用的结果（这个结果也必须是可以调用的）来调用正在定义的函数。不管是哪种情况，最后一次或唯一一次调用的返回值都会绑定到定义的函数名称上。

Answer 2

如果你想给你的装饰器传递参数，你需要始终把它当作一个函数来调用：

@d()
def func():
    pass

否则，你就得尝试去判断参数的不同，也就是说，你需要神奇地猜测调用者的意思。不要创建一个需要猜测的接口；一开始就清楚地表达你的意思。

换句话说，一个函数要么是装饰器，要么是装饰器工厂；它不应该同时是两者。

注意，如果你只是想存储一个值，其实你不需要写一个类。

def d(msg='my default message'):
    def decorator(func):
        def newfn():
            print msg
            return func()
        return newfn
    return decorator

@d('This is working')
def hello():
    print 'hello world !'

@d()
def hello2():
    print 'also hello world'

Answer 3

我找到一个例子，你可以使用 @trace 或者 @trace('msg1','msg2')：挺不错的！

def trace(*args):
    def _trace(func):
        def wrapper(*args, **kwargs):
            print enter_string
            func(*args, **kwargs)
            print exit_string
        return wrapper
    if len(args) == 1 and callable(args[0]):
        # No arguments, this is the decorator
        # Set default values for the arguments
        enter_string = 'entering'
        exit_string = 'exiting'
        return _trace(args[0])
    else:
        # This is just returning the decorator
        enter_string, exit_string = args
        return _trace