How to build a decorator with optional parameters?

Question

I would like to make a decorator which could be used with or without a parameter : Something like this :

class d(object):
    def __init__(self,msg='my default message'):
        self.msg = msg
    def __call__(self,fn):
        def newfn():
            print self.msg
            return fn()
        return newfn

@d('This is working')
def hello():
    print 'hello world !'

@d
def too_bad():
    print 'does not work'

In my code, only the use of decorator with parameter is working: How to proceed to have both working (with and without parameter)?

Solution 1

I found an example, you can use @trace or @trace('msg1','msg2'): nice!

def trace(*args):
    def _trace(func):
        def wrapper(*args, **kwargs):
            print enter_string
            func(*args, **kwargs)
            print exit_string
        return wrapper
    if len(args) == 1 and callable(args[0]):
        # No arguments, this is the decorator
        # Set default values for the arguments
        enter_string = 'entering'
        exit_string = 'exiting'
        return _trace(args[0])
    else:
        # This is just returning the decorator
        enter_string, exit_string = args
        return _trace

Solution 2

If you want to take parameters to your decorator, you need to always call it as a function:

@d()
def func():
    pass

Otherwise, you need to try to detect the difference in parameters--in other words, you need to magically guess what the caller means. Don't create an API that needs to guess; consistently say what you mean to begin with.

In other words, a function should either be a decorator, or a decorator factory; it shouldn't be both.

Note that if all you want to do is store a value, you don't need to write a class.

def d(msg='my default message'):
    def decorator(func):
        def newfn():
            print msg
            return func()
        return newfn
    return decorator

@d('This is working')
def hello():
    print 'hello world !'

@d()
def hello2():
    print 'also hello world'

Solution 3

If you don't mind relying on using named arguments, I made something similar to what you need:

def cached_property(method=None, get_attribute=lambda a: '_%s_cached' % (a,)):
    """
    Caches an object's attribute.

    Can be used in the following forms:
    @cached_property
    @cached_property()
    @cached_property(get_attribute=lambda x: 'bla')

    @param method: the method to memoizes
    @param get_attribute: a callable that should return the cached attribute
    @return a cached method
    """
    def decorator(method):
        def wrap(self):
            private_attribute = get_attribute(method.__name__)
            try:
                return getattr(self, private_attribute)
            except AttributeError:
                setattr(self, private_attribute, method(self))
                return getattr(self, private_attribute)
        return property(wrap)
    if method:
        # This was an actual decorator call, ex: @cached_property
        return decorator(method)
    else:
        # This is a factory call, ex: @cached_property()
        return decorator

This works because only one non keyword argument, the function decorated is passed to the decorator.

Notice that I also used the arguments passed to the decorated function, in this case 'self'.

Solution 4

This would work.

def d(arg):
    if callable(arg):  # Assumes optional argument isn't.
        def newfn():
            print('my default message')
            return arg()
        return newfn
    else:
        def d2(fn):
            def newfn():
                print(arg)
                return fn()
            return newfn
        return d2

@d('This is working')
def hello():
    print('hello world !')

@d  # No explicit arguments will result in default message.
def hello2():
    print('hello2 world !')

@d('Applying it twice')
@d('Would also work')
def hello3():
    print('hello3 world !')

hello()
hello2()
hello3()

Output:

This is working
hello world !
my default message
hello2 world !
Applying it twice
Would also work
hello3 world !

If a decorator function @invocation isn't passed any explicit arguments, it is called with the function defined in the following def. If it is passed arguments, then it is first called with them and then the result of that preliminary call (which must itself also be a callable) is called with the function being defined. Either way, the return value of the last or only call is bound to the defined function name.

Solution 5

You have to detect if the argument to the decorator is a function, and use a simple decorator in that case. And then you need to hope that you never need to pass only a function to the parametrized decorator.

View more solutions

Comments

[Muhammad Alkarouri]

This is probably a good advice, but it is not correct that a function can't do both, as Ignacio Vazquez-Abrams explains. It is probably better to explain that in the answer.

[Glenn Maynard]

@Muhammad: I didn't say it can't, I said it shouldn't.

[Muhammad Alkarouri]

I understand. But the value of the answer would be higher if this point is explained a bit better. Just saying.

[ksrini]

I guess it works for your case and would for similar cases. But what if the decorator's argument is actually a single callable? How would you differentiate between the decorated function and the argument?

[martineau]

@ksrini: Ignacio pointed that out in his answer years ago.

[csiu]

You can also get around this by using keyword arguments (such as @trace(default=...)).

[Andy]

This is a much simpler example to understand than the truly optional one provided by Eric. Python newbies might get lost in that one quicker than they would in this one.

[martineau]

I don't consider checking the type or number of arguments passed as "guessing", but do think that having to remember and use the decorator via @d() instead of usual @d in some case poor advice—including because it looks very weird/odd.

[adkl]

Excellent, what an answer!

[Cecil Curry]

No magical guesswork is required. Violating Python conventions for no good reason is patently bad advice. Determining whether or not a decorator function is being called as a decorator or decorator factory trivially reduces to accepting only optional keyword rather than positional arguments (e.g., replacing *args with **kwargs in Eric's rightfully accepted solution). No ambiguity exists; ergo, there exists no reason to invent new calling conventions.

[Cecil Curry]

To avoid both this ambiguity and one unnecessary level of inner closure nesting, see ForeverWintr's single-line improvement of PatBenavente's similar answer elsewhere.

[stenci]

This is wrong: "you need to always call it as a function". Perhaps you wanted to say "should" instead of "need"?