Search code examples
pythonpython-2.7classdecoratorpython-decorators

Decorating class methods to build a method registry


What I want to achieve:

I want to decorate class method definitions such that the class can build a registry of of certain methods. This was intended to work as follows:

class NForm(FormFiller):
    @fills('name')
    def get_name(self):
        return 'gets name'

class NAForm(NForm):
    @fills('address')
    def get_address(self):
        return 'gets address'

class NBForm(NForm):
    @fills('birthday')
    def get_birthday(self):
        return 'gets birthday'

na = NAForm().filled_form() == {'address': 'gets address',
                                'name': 'gets name'}

nb = NBForm().filled_form() == {'birthday': 'gets birthday',
                                'name': 'gets name'}

What I attempted

It seemed easy enough, so I wrote this code for the parent class.

class FormFiller(object):
    _fills = {}

    @classmethod
    def fills(cls, field_name):
        def wrapper(func):
            cls._fills[field_name] = func.__name__
            return func
        return wrapper

    def filled_form(self):
        return {field_name: getattr(self, func)() for field_name, func in self._fills.items()}

And replaced the decorators above with e.g. @NAForm.fills('address'). However, this is of course not possible as NAForm is not yet defined within its own definition. I can only ever write into any of parents' _fills attribute.

It feels like the desired use case should be possible, but I currently have no idea how to achieve this. If I wanted to manually achieve the behaviour above, I could add an in-between class for every inheritance, e.g.

class FormFiller(object):
    _fills = {}
    @classmethod
    def fills(cls, field_name):
        print 'in', cls, 'adding', field_name
        def wrapper(func):
            # make a copy since _fills is like a mutable default
            cls._fills = dict(cls._fills)
            # update
            cls._fills[field_name] = func.__name__
            return func
        return wrapper

    def filled_form(self):
        return {field_name: getattr(self, func)() for field_name, func in self._fills.items()}


class NFormBase(FormFiller):
    pass

class NForm(NFormBase):
    @NFormBase.fills('name')
    def get_name(self):
        return 'gets name'

class NAFormBase(NForm):
    pass

class NAForm(NAFormBase):
    @NAFormBase.fills('address')
    def get_address(self):
        return 'gets address'

class NBFormBase(NForm):
    pass

class NBForm(NBFormBase):
    @NBFormBase.fills('birthday')
    def get_age(self):
        return 'gets birthday'


print FormFiller().filled_form()  # == {}
print NForm().filled_form()  # == {'name': 'gets name'}
print NAForm().filled_form()  # == {'name': 'gets name', 'address': 'gets address'}
print NBForm().filled_form()  # == {'birthday': 'gets birthday', 'name': 'gets name'}

This seems to work, but a) requires you to add an inbetween class for every inheritance step and b) copies the _fills dictionary much more frequently than is necessary (on every decoration instead of once per class creation).

Is there a better way to do this? I'd be grateful for any pointers in the right direction. Are metaclasses what I'm looking for?

Thanks!


Solution

  • It turns out metaclasses are what I needed, I had never used them before and I only just found the appropriate question on SO: Auto-register class methods using decorator

    registry = {}
    
    class RegisteringType(type):
        def __init__(cls, name, bases, attrs):
            registry[name] = {}
    
            for base in bases:
                registry[name].update(**registry.get(base.__name__, {}))
    
            for key, val in attrs.iteritems():
                properties = getattr(val, 'register', None)
                if properties is not None:
                    registry[name][key] = properties
    
    def register(*args):
        def decorator(f):
            f.register = tuple(args)
            return f
        return decorator
    
    class FormFiller(object):
        __metaclass__ = RegisteringType
    
        def filled_form(self):
            print type(self).__name__
            return registry[type(self).__name__]
    
    class NForm(FormFiller):
        @register('name')
        def get_name(self):
            return 'gets name'
    
    class NAForm(NForm):
        @register('address')
        def get_address(self):
            return 'gets address'
    
    class NBForm(NForm):
        @register('birthday')
        def get_age(self):
            return 'gets birthday'