Diego Búrigo Zacarão’s Weblog

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.