How to create a Python decorator that can be used either with or without parameters?
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
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():
...
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.