How to create a Python 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

Answer 1

Since no one mentioned this, there is also a solution utilizing callable class which I find more elegant, especially in cases where the decorator is complex and one may wish to split it to multiple methods(functions). This solution utilizes __new__ magic method to do essentially what others have pointed out. First detect how the decorator was used than adjust return appropriately.

class decorator_with_arguments(object):

    def __new__(cls, decorated_function=None, **kwargs):

        self = super().__new__(cls)
        self._init(**kwargs)

        if not decorated_function:
            return self
        else:
            return self.__call__(decorated_function)

    def _init(self, arg1="default", arg2="default", arg3="default"):
        self.arg1 = arg1
        self.arg2 = arg2
        self.arg3 = arg3

    def __call__(self, decorated_function):

        def wrapped_f(*args):
            print("Decorator arguments:", self.arg1, self.arg2, self.arg3)
            print("decorated_function arguments:", *args)
            decorated_function(*args)

        return wrapped_f

@decorator_with_arguments(arg1=5)
def sayHello(a1, a2, a3, a4):
    print('sayHello arguments:', a1, a2, a3, a4)

@decorator_with_arguments()
def sayHello(a1, a2, a3, a4):
    print('sayHello arguments:', a1, a2, a3, a4)

@decorator_with_arguments
def sayHello(a1, a2, a3, a4):
    print('sayHello arguments:', a1, a2, a3, a4)

If the decorator is used with arguments, than this equals:

result = decorator_with_arguments(arg1=5)(sayHello)(a1, a2, a3, a4)

One can see that the arguments arg1 are correctly passed to the constructor and the decorated function is passed to __call__

But if the decorator is used without arguments, than this equals:

result = decorator_with_arguments(sayHello)(a1, a2, a3, a4)

You see that in this case the decorated function is passed directly to the constructor and call to __call__ is entirely omitted. That is why we need to employ logic to take care of this case in __new__ magic method.

Why can't we use __init__ instead of __new__? The reason is simple: python prohibits returning any other values than None from __init__

WARNING

This approcach has one side effect. It will not preserve function signature!

Answer 2

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 :-)

Answer 3

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

Answer 4

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.

Answer 5

In fact, the caveat case in @bj0's solution can be checked easily:

def meta_wrap(decor):
    @functools.wraps(decor)
    def new_decor(*args, **kwargs):
        if len(args) == 1 and len(kwargs) == 0 and callable(args[0]):
            # this is the double-decorated f. 
            # Its first argument should not be a callable
            doubled_f = decor(args[0])
            @functools.wraps(doubled_f)
            def checked_doubled_f(*f_args, **f_kwargs):
                if callable(f_args[0]):
                    raise ValueError('meta_wrap failure: '
                                'first positional argument cannot be callable.')
                return doubled_f(*f_args, **f_kwargs)
            return checked_doubled_f 
        else:
            # decorator arguments
            return lambda real_f: decor(real_f, *args, **kwargs)

    return new_decor

Here are a few test cases for this fail-safe version of meta_wrap.

    @meta_wrap
    def baddecor(f, caller=lambda x: -1*x):
        @functools.wraps(f)
        def _f(*args, **kwargs):
            return caller(f(args[0]))
        return _f

    @baddecor  # used without arg: no problem
    def f_call1(x):
        return x + 1
    assert f_call1(5) == -6

    @baddecor(lambda x : 2*x) # bad case
    def f_call2(x):
        return x + 1
    f_call2(5)  # raises ValueError

    # explicit keyword: no problem
    @baddecor(caller=lambda x : 100*x)
    def f_call3(x):
        return x + 1
    assert f_call3(5) == 600

Answer 6

To give a more complete answer than the above:

"Is there a way to build a decorator that can be used both with and without arguments ?"

No there is no generic way because there is currently something missing in the python language to detect the two different use cases.

However Yes as already pointed out by other answers such as bj0s, there is a clunky workaround that is to check the type and value of the first positional argument received (and to check if no other arguments have non-default value). If you are guaranteed that users will never pass a callable as first argument of your decorator, then you can use this workaround. Note that this is the same for class decorators (replace callable by class in the previous sentence).

To be sure of the above, I did quite a bit of research out there and even implemented a library named decopatch that uses a combination of all strategies cited above (and many more, including introspection) to perform "whatever is the most intelligent workaround" depending on your need.

But frankly the best would be not to need any library here and to get that feature straight from the python language. If, like myself, you think that it is a pity that the python language is not as of today capable of providing a neat answer to this question, do not hesitate to support this idea in the python bugtracker: https://bugs.python.org/issue36553 !

Thanks a lot for your help making python a better language :)