I’ve just been doing something that probably should never be done in Python (again), and I figured it might be useful to record this little snippet for posterity. Many Python programmers are probably familiar with decorators, which modify functions by wrapping them with other functions. For example, if you want to print ‘entering’ and ‘exiting’ to trace when a function’s execution begins and ends, you could do that with a decorator:

def trace(func):
    def wrapper(*args, **kwargs):
        print 'entering'
        func(*args, **kwargs)
        print 'exiting'
    return wrapper

@trace
def f(a,b):
     print 'in f: ' + str(a) + ',' + str(b)

Note that in practice there are better ways to do this, but it’s just an example.

Anyway, it’s also possible to customize the behavior of decorator itself by providing arguments, for instance if you wanted to print different strings instead of ‘entering’ and ‘exiting’. But in that case, your “decorator” function is actually a wrapper that creates the real decorator and returns it:

def trace(enter_string, exit_string):
    def _trace(func):
        def wrapper(*args, **kwargs):
            print enter_string
            func(*args, **kwargs)
            print exit_string
        return wrapper
    return _trace

@trace('calling f', 'returning')
def f(a,b):
     print 'in f: ' + str(a) + ',' + str(b)

OK, so far so good. Slightly mind-bending, but that’s Python for you.

Now, what if you want both ways to work? What if you want a decorator that can be used with or without arguments, as either

@trace
def f(a,b):
     print 'in f: ' + str(a) + ',' + str(b)

or

@trace('calling f', 'returning')
def f(a,b):
     print 'in f: ' + str(a) + ',' + str(b)

Well, in the former case the decorator will be called with one argument, the function to be decorated. You can identify this case in an if statement and treat it separately. Like so:

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

In this example, _trace is the “real” decorator, the thing that gets called to wrap f. The top-level function trace either stores its arguments and returns the decorator itself (if it has arguments), or calls the decorator to do its job and returns the result (if there are no arguments provided, other than the function to be wrapped of course). Obviously, this won’t work if you have a decorator that needs to take one argument (not the function it’s wrapping) which is itself a callable function — in that case you’re on your own.

In fact — if you’re not confused yet — you can even make a decorator that turns your decorator functions into optional-argument decorators. A meta-decorator, I guess. Here’s one way to do it:

def opt_arguments(func):
    def meta_wrapper(*args, **kwargs):
        if len(args) == 1 and callable(args[0]):
            # No arguments, this is the decorator
            # Set default values for the arguments
            return func(args[0])
        else:
            def meta_func(inner_func):
                return func(inner_func, *args, **kwargs)
            return meta_func
    return meta_wrapper

You can use this to decorate a decorator function which has one required argument and default values for all other arguments:

@opt_arguments
def trace(func, enter_string='entering', exit_string='exiting'):
    def wrapper(*args, **kwargs):
        print enter_string
        func(*args, **kwargs)
        print exit_string
    return wrapper

@trace
def f(a,b):
    print 'in f(%s,%s)' % (a,b)

@trace('call g', 'return g')
def g(a,b):
    print 'in g(%s,%s)' % (a,b)

By the way, I haven’t tried this when the decorator is a class or callable object instead of a function, but I think it should work.