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

Question

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.

Solution 1

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

Solution 2

Using keyword arguments with default values (as suggested by kquinn) is a good idea, but will require you to include the parenthesis:

@redirect_output()
def foo():
    ...

If you would like a version that works without the parenthesis on the decorator you will have to account both scenarios in your decorator code.

If you were using Python 3.0 you could use keyword only arguments for this:

def redirect_output(fn=None,*,destination=None):
  destination = sys.stderr if destination is None else destination
  def wrapper(*args, **kwargs):
    ... # your code here
  if fn is None:
    def decorator(fn):
      return functools.update_wrapper(wrapper, fn)
    return decorator
  else:
    return functools.update_wrapper(wrapper, fn)

In Python 2.x this can be emulated with varargs tricks:

def redirected_output(*fn,**options):
  destination = options.pop('destination', sys.stderr)
  if options:
    raise TypeError("unsupported keyword arguments: %s" % 
                    ",".join(options.keys()))
  def wrapper(*args, **kwargs):
    ... # your code here
  if fn:
    return functools.update_wrapper(wrapper, fn[0])
  else:
    def decorator(fn):
      return functools.update_wrapper(wrapper, fn)
    return decorator

Any of these versions would allow you to write code like this:

@redirected_output
def foo():
    ...

@redirected_output(destination="somewhere.log")
def bar():
    ...

Solution 3

I know this is an old question, but I really don't like any of the techniques proposed so I wanted to add another method. I saw that django uses a really clean method in their login_required decorator in django.contrib.auth.decorators. As you can see in the decorator's docs, it can be used alone as @login_required or with arguments, @login_required(redirect_field_name='my_redirect_field').

The way they do it is quite simple. They add a kwarg (function=None) before their decorator arguments. If the decorator is used alone, function will be the actual function it is decorating, whereas if it is called with arguments, function will be None.

Example:

from functools import wraps

def custom_decorator(function=None, some_arg=None, some_other_arg=None):
    def actual_decorator(f):
        @wraps(f)
        def wrapper(*args, **kwargs):
            # Do stuff with args here...
            if some_arg:
                print(some_arg)
            if some_other_arg:
                print(some_other_arg)
            return f(*args, **kwargs)
        return wrapper
    if function:
        return actual_decorator(function)
    return actual_decorator

---

@custom_decorator
def test1():
    print('test1')

>>> test1()
test1

---

@custom_decorator(some_arg='hello')
def test2():
    print('test2')

>>> test2()
hello
test2

---

@custom_decorator(some_arg='hello', some_other_arg='world')
def test3():
    print('test3')

>>> test3()
hello
world
test3

I find this approach that django uses to be more elegant and easier to understand than any of the other techniques proposed here.

Solution 4

Several answers here already address your problem nicely. With respect to style, however, I prefer solving this decorator predicament using functools.partial, as suggested in David Beazley's Python Cookbook 3:

from functools import partial, wraps

def decorator(func=None, foo='spam'):
    if func is None:
         return partial(decorator, foo=foo)

    @wraps(func)
    def wrapper(*args, **kwargs):
        # do something with `func` and `foo`, if you're so inclined
        pass

    return wrapper

While yes, you can just do

@decorator()
def f(*args, **kwargs):
    pass

without funky workarounds, I find it strange looking, and I like having the option of simply decorating with @decorator.

As for the secondary mission objective, redirecting a function's output is addressed in this Stack Overflow post.

---

If you want to dive deeper, check out Chapter 9 (Metaprogramming) in Python Cookbook 3, which is freely available to be read online.

Some of that material is live demoed (plus more!) in Beazley's awesome YouTube video Python 3 Metaprogramming.

Happy coding :)

Solution 5

You need to detect both cases, for example using the type of the first argument, and accordingly return either the wrapper (when used without parameter) or a decorator (when used with arguments).

from functools import wraps
import inspect

def redirect_output(fn_or_output):
    def decorator(fn):
        @wraps(fn)
        def wrapper(*args, **args):
            # Redirect output
            try:
                return fn(*args, **args)
            finally:
                # Restore output
        return wrapper

    if inspect.isfunction(fn_or_output):
        # Called with no parameter
        return decorator(fn_or_output)
    else:
        # Called with a parameter
        return decorator

When using the @redirect_output("output.log") syntax, redirect_output is called with a single argument "output.log", and it must return a decorator accepting the function to be decorated as an argument. When used as @redirect_output, it is called directly with the function to be decorated as an argument.

Or in other words: the @ syntax must be followed by an expression whose result is a function accepting a function to be decorated as its sole argument, and returning the decorated function. The expression itself can be a function call, which is the case with @redirect_output("output.log"). Convoluted, but true :-)

View more solutions

Comments

[user1066101]

The default-looking @redirect_output is remarkably uninformative. I'd suggest that it's a bad idea. Use the first form and simplify your life a lot.

[rog]

interesting question though - until i saw it and looked through the documentation, i'd have assumed that @f was the same as @f(), and i still think it should be, to be honest (any provided arguments would just be tacked on to the function argument)

[Tomasz Gandor]

This decorator-factory/decorator pattern is nice, with first default argument function=None, I would go further, and make the remaining arguments thereafter keyword-only.

[sth]

If you say @dec(abc) the function is not passed directly to dec. dec(abc) returns something, and this return value is used as the decorator. So dec(abc) has to return a function, which then gets the decorated function passed as an parameter. (Also see thobes code)

[ehabkost]

this can't be used as a decorator like in the @redirect_output("somewhere.log") def foo() example in the question.

[lum]

What do you put in your code here? How do you call the function that is decorated? fn(*args, **kwargs) doesn't work.

[Pragy Agarwal]

Thanks. This is helpful!

[Omer Ben Haim]

i think there is a much simpler answer, create a class which will the the decorator with optional arguments. create another function with the same arguments with defaults and return a new instance of the decorator classes. should look something like: def f(a = 5): return MyDecorator( a = a) and class MyDecorator( object ): def __init__( self, a = 5 ): .... sorry its hard writing it in a comment but i hope this is simple enough to understand

[Dustin Wyatt]

Yeah, I like this method. Do note that you have to use kwargs when calling the decorator otherwise the first positional arg is assigned to function and then things break because the decorator tries to call that first positional arg as if it were your decorated function.

[Russell Smith]

"than the above" isn't a useful phrase in stackoverflow. Different people will see answers in different orders over time. It's impossible to know what answer you're referring to.

[smarie]

Thanks for spotting this @BryanOakley. Indeed, if you had found it useful and voted it up there would have been less messages above. I edited the message accordingly

[Tomasz Gandor]

Yes, the first argument is not a kwarg, it's a positional argument with a default. But you could make the rest of arguments keyword-only.

[Noam Nol]

It will not work if the first parameter is a class. Instead of callable(args[0]), you can detect classes with: isinstance(args[0], types.FunctionType)

[bj0]

I mentioned that that was one of the assumptions in the example I provided. For special cases like passing a class, you would need to modify it to fit your case.

[n8henrie]

I like this style the best, thanks for sharing. For others, I will note that if you try to mutate foo within wrapper, you may get an UnboundLocalError, in which case you should declare nonlocal foo (or perhaps choose a different local variable name, bar, and set bar = foo). See also: stackoverflow.com/a/57184656/1588795

[Björn Pollex]

Note that this solution assumes that calling code always uses keyword arguments. In this case @decorator('foo') would not behave as expected.