Generic Python Decorators
Aug 1st, 2011 by diegobz
The past days I came across a need to check specific attributes in a class instance before calling a specific method of it. I decided to use decorators, but my problem was that I needed to check different attributes and I would like to do it without duplicating code and also in an elegant way. So I started to check how I could create decorators in a generic way, to be reusable, something that would end up as the following:
class Obj(object):
@need_foo
def foo(self):
return self._foo
@need_bar
def bar(self):
return self._bar
o=Obj()
o.foo()
AttributeError: Attr _foo not specified.
Then it turned out that I could have two possibilities. The two implementations basically do exactly the same thing, the only difference is that one is simply a function and the other one is a class based decorator that have the flexibility of all resources of a class over a function. It doesn’t mean you will always use it, but it gives you the power.
class CheckAttr(object):
"""
Class decorator to help checking whether an attribute of an
object is set before calling the wanted (decorated) method.
This class decorator must be instantiated with a 'attr' and
'message', which are respectively the attr name to be
checked and the exception message in case it was not set
yet.
"""
def __init__(self, attr, message):
self.attr = attr
self.message = message
def __call__(self, func):
def _wrapper(*args, **kwargs):
if not getattr(args[0], self.attr, None):
raise AttributeError(self.message)
return func(*args, **kwargs)
return _wrapper
def check_attr(attr, message):
"""
Decorator to help checking whether an attribute of an
object is set before calling the wanted (decorated)
method of an object.
This decorator must be called with a 'attr' and 'message',
which are respectively the attribute name to be checked
and the exception message in case it was not set yet.
"""
def _check_func(func):
def _wrapper(*args, **kwargs):
if not getattr(args[0], attr, None):
raise AttributeError(message)
return func(*args, **kwargs)
return _wrapper
return _check_func
#
# DECORATOR WRAPPERS #
#
def need_bar(func):
return CheckAttr('_bar', 'Attr _bar not specified.')(func)
def need_qux(func):
return check_attr('_qux', 'Attr _qux not specified.')(func)
The code above allow us to use those two kinds of decorators in the following ways. Examples on how to invoke the code are in the docstrings as doctest:
class Obj(object):
"""
>>> o = Obj()
>>> o.foo()
Traceback (most recent call last):
...
AttributeError: Attr _foo not specified. Use `bind_foo`.
>>> o.bar()
Traceback (most recent call last):
...
AttributeError: Attr _bar not specified.
>>> o.baz()
Traceback (most recent call last):
...
AttributeError: Attr _baz not specified. Use `bind_baz`.
>>> o.qux()
Traceback (most recent call last):
...
AttributeError: Attr _qux not specified.
"""
# Calling the class based decorator directly
@CheckAttr('_foo', 'Attr _foo not specified. Use `bind_foo`.')
def foo(self):
return self._foo
def bind_foo(self, foo):
self._foo = foo
# Calling the class based decorator using a wrapper
@need_bar
def bar(self):
return self._bar
# Calling the function based decorator directly
@check_attr('_baz', 'Attr _baz not specified. Use `bind_baz`.')
def baz(self):
return self._baz
def bind_baz(self, baz):
self._baz = baz
# Calling the function based decorator using a wrapper
@need_qux
def qux(self):
return self._qux
In the end I used the function based decorator with the wrappers (@need_attribute), because it was enough for what I wanted. But it’s quite obvious now to me that decorators can be used in much more powerful ways.
