A sqlalchemy-continuum plugin for making model methods and properties available on the versioned model

method_proxy.py

""" A sqlalchemy-continuum plugin for making model methods
    and properties available on the versioned model
"""

from functools import wraps
from sqlalchemy_continuum.plugins import Plugin


def override_in_versioned(function=None):
    """ Use the explicitly mark methods which should be
        available on the Versioned model (even if this
        means the get overridden)
    """
    def decorator(f):
        # No wrapping required, just set a flag on the
        # function which we can use later
        f._override_in_versioned = True
        return f
    if function:
        return decorator(function)
    else:
        return decorator


class MethodProxyPlugin(Plugin):

    def make_accessor(self):
        def accessor(version_obj, name):
            # Only try this if we have a parent
            if version_obj.version_parent:
                # Get the attribute we are after (or raise an AttributeError)
                attr = getattr(version_obj.version_parent, name)
                if callable(attr):
                    # If it is callable we need to make sure we pass the versioned
                    # object as 'self', rather than the parent object
                    @wraps(attr)
                    def custom_self(*args, **kwargs):
                        return attr.im_func(version_obj, *args, **kwargs)
                    return custom_self
                else:
                    # Otherwise simply return the attribute
                    return attr
            else:
                raise AttributeError("%s has no attribute '%s'" % (repr(version_obj), repr(name)))
        return accessor

    def after_version_class_built(self, parent_cls, version_cls):
        # General accessor for resolving attributes not
        # found on the versioned class
        setattr(version_cls, '__getattr__', self.make_accessor())

        # Now override any methods if it was specifically requested via the
        # @override_in_versioned decorator
        for attr_name in dir(parent_cls):
            try:
                attr_value = getattr(parent_cls, attr_name)
            except:
                # The above causes hybrid properties to be called, which can cause
                # errors. A better solution to this would be great.
                continue

            # Do the overriding
            if getattr(attr_value, '_override_in_versioned', False):
                setattr(version_cls, attr_name, attr_value.im_func)

How would you use this?