How to build a decorator with optional parameters?

Eric picture Eric · Oct 14, 2010 · Viewed 25.2k times · Source

I would like to make a decorator which could be used with or without a parameter : Something like this :

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'

In my code, only the use of decorator with parameter is working: How to proceed to have both working (with and without parameter)?

Answer

Eric picture Eric · Oct 14, 2010

I found an example, you can use @trace or @trace('msg1','msg2'): nice!

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