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 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
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
OTHER TIPS
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():
...
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 :-)
A python decorator is called in a fundamentally different way depending on whether you give it arguments or not. The decoration is actually just a (syntactically restricted) expression.
In your first example:
@redirect_output("somewhere.log")
def foo():
....
the function redirect_output
is called with the
given argument, which is expected to return a decorator
function, which itself is called with foo
as an argument,
which (finally!) is expected to return the final decorated function.
The equivalent code looks like this:
def foo():
....
d = redirect_output("somewhere.log")
foo = d(foo)
The equivalent code for your second example looks like:
def foo():
....
d = redirect_output
foo = d(foo)
So you can do what you'd like but not in a totally seamless way:
import types
def redirect_output(arg):
def decorator(file, f):
def df(*args, **kwargs):
print 'redirecting to ', file
return f(*args, **kwargs)
return df
if type(arg) is types.FunctionType:
return decorator(sys.stderr, arg)
return lambda f: decorator(arg, f)
This should be ok unless you wish to use a function as an argument to your decorator, in which case the decorator will wrongly assume it has no arguments. It will also fail if this decoration is applied to another decoration that does not return a function type.
An alternative method is just to require that the decorator function is always called, even if it is with no arguments. In this case, your second example would look like this:
@redirect_output()
def foo():
....
The decorator function code would look like this:
def redirect_output(file = sys.stderr):
def decorator(file, f):
def df(*args, **kwargs):
print 'redirecting to ', file
return f(*args, **kwargs)
return df
return lambda f: decorator(file, f)
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.
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 :)
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
Have you tried keyword arguments with default values? Something like
def decorate_something(foo=bar, baz=quux):
pass
Generally you can give default arguments in Python...
def redirect_output(fn, output = stderr):
# whatever
Not sure if that works with decorators as well, though. I don't know of any reason why it wouldn't.
Building on vartec's answer:
imports sys
def redirect_output(func, output=None):
if output is None:
output = sys.stderr
if isinstance(output, basestring):
output = open(output, 'w') # etc...
# everything else...