How to build a decorator with optional parameters?

Viewed 32080

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

5 Answers

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

Related