How to create a Python decorator that can be used either with or without parameters?

elifiner picture elifiner · Mar 17, 2009 · Viewed 18.1k times · Source

I'd like to create a Python decorator that can be used either with parameters:

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

or without them (for instance to redirect the output to stderr by default):

@redirect_output
def foo():
    ....

Is that at all possible?

Note that I'm not looking for a different solution to the problem of redirecting output, it's just an example of the syntax I'd like to achieve.

Answer

bj0 picture bj0 · Jan 19, 2013

I know this question is old, but some of the comments are new, and while all of the viable solutions are essentially the same, most of them aren't very clean or easy to read.

Like thobe's answer says, the only way to handle both cases is to check for both scenarios. The easiest way is simply to check to see if there is a single argument and it is callabe (NOTE: extra checks will be necessary if your decorator only takes 1 argument and it happens to be a callable object):

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)

In the first case, you do what any normal decorator does, return a modified or wrapped version of the passed in function.

In the second case, you return a 'new' decorator that somehow uses the information passed in with *args, **kwargs.

This is fine and all, but having to write it out for every decorator you make can be pretty annoying and not as clean. Instead, it would be nice to be able to automagically modify our decorators without having to re-write them... but that's what decorators are for!

Using the following decorator decorator, we can deocrate our decorators so that they can be used with or without arguments:

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

Now, we can decorate our decorators with @doublewrap, and they will work with and without arguments, with one caveat:

I noted above but should repeat here, the check in this decorator makes an assumption about the arguments that a decorator can receive (namely that it can't receive a single, callable argument). Since we are making it applicable to any generator now, it needs to be kept in mind, or modified if it will be contradicted.

The following demonstrates its use:

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