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

Viewed 25452

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.

14 Answers

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)

To complete the other answers:

"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. It comes bundled with two modes: nested and flat.

In "nested mode" you always return a function

from decopatch import function_decorator

@function_decorator
def add_tag(tag='hi!'):
    """
    Example decorator to add a 'tag' attribute to a function. 
    :param tag: the 'tag' value to set on the decorated function (default 'hi!).
    """
    def _apply_decorator(f):
        """
        This is the method that will be called when `@add_tag` is used on a 
        function `f`. It should return a replacement for `f`.
        """
        setattr(f, 'tag', tag)
        return f
    return _apply_decorator

while in "flat mode" your method is directly the code that will be executed when the decorator is applied. It is injected with the decorated function object f:

from decopatch import function_decorator, DECORATED

@function_decorator
def add_tag(tag='hi!', f=DECORATED):
    """
    Example decorator to add a 'tag' attribute to a function.
    :param tag: the 'tag' value to set on the decorated function (default 'hi!).
    """
    setattr(f, 'tag', tag)
    return f

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

This does the job without no fuss:

from functools import wraps

def memoize(fn=None, hours=48.0):
  def deco(fn):
    @wraps(fn)
    def wrapper(*args, **kwargs):
      return fn(*args, **kwargs)
    return wrapper

  if callable(fn): return deco(fn)
  return deco

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!

this work for me:

def redirect_output(func=None, /, *, output_log='./output.log'):
    def out_wrapper(func):
        def wrapper(*args, **kwargs):
            res = func(*args, **kwargs)
            print(f"{func.__name__} finished, output_log:{output_log}")
            return res

        return wrapper

    if func is None:
        return out_wrapper  # @redirect_output()
    return out_wrapper(func)  # @redirect_output


@redirect_output
def test1():
    print("running test 1")


@redirect_output(output_log="new.log")
def test2():
    print("running test 2")

test1()
print('-----')
test2()

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...
Related