Python - reload a module and update the methods for pre-existing versions of its class(es).

reloadd.py

# -*- coding: utf-8 -*-
import sys
from types import ModuleType, ClassType

def can_reload_methods(klass):
    klass.__CAN_RELOAD_METHODS__ = True
    return klass
    

def reloadd(m):
    """
        ``m`` can be a string or an actual module
        
        Any classes in the module that are decorated with ``@can_reload_methods``
        will have their old class methods updated to the new version, meaning 
        that any existing instances of the class will now be using the new methods.
        
        This is not a silver bullet at all, but in many cases it will be useful.
        
        If the module is not in the module cache we just don't do anything here.
        
        Example
        -----
        
        Imagine we have a class ``someClass`` in a module ``someModule``. First
        we create an instance, ``a``, of the class and call one of its functions::

            import someModule
            a = someModule.someClass()
            a.something()
            >>> "this is something version 1"
            
        Next, suppose we go and edit the ``something`` method in ``someModule.py``,
        lets say it should now print ``"this is version 2"``.
        
        Having done that, if we call ``a.something()`` again we will be dissapointed to find that
        the result hasn't changed: we still get ``this is soemthing version  1``.
        
        You might try doing ``reload(someModule)`` but it wont help.   Which is 
        where this special ``reloadd`` function comes in handy::
        
            reloadd(someModule)
            a.something()
            >>> "this is version 2"
        
        Which is what we wanted.  This is not bullet proof reloading by any means
        but in most cases it should update existing instances of ``someModule`` and
        any instances of classes that use ``someModule`` as a base class.
        
        Note that ``reloadd`` only applies to the single file you specify, it
        does not do any recursion whatsoever...but in most cases that will be
        the behaviour you actually wanted...this can all get seriously confusing.        
        
        Note that for a class to be recognised by ``reloadd`` it needs to be
        decorated with the ``@can_reload_methods`` decorator.  This isn't actually
        a neccessary part of the implementation, but it seems fairly sensible.
        If that's annoying, you could edit this function to reload all classes 
        within the requested module, or have an optional flag.
        (If you add the ``@can_reload_methods`` decorator after creating instances
        and then try and ``reloadd``, the methods will be applied because the
        new version of the class has the decorator even though the old version does not.)
        
        How it works
        ------

        The main idea is to look in the module cache to find the old version of
        the relevant class(es), then do a standard reload, and then overwrite/add
        the new methods to the old class.
        
        This more or less works, but there are some additional complications, the
        most significant of which is that if we want to do repeated ``reloadd`` calls
        we need to be able to chain back to the oldest version of the class, because
        that is the version which has the active instances (well, there may well be
        instances of a variety of different versions of the class and we have to
        update them all.)
        
        Warning
        --------
        
        There are no doubt many unforseen ways this could go wrong.

        Methods and properties are updated.  
                
        Old methods that have been removed are not deleted, i.e.
        only methods named in the new version are updated. 
        
        Not sure about staicmethods, classmethods and methods wrapped with
        other decorators.
        
    """
    if isinstance(m,ModuleType) is False and m not in sys.modules: 
        return
        
    name_m = m if isinstance(m,str) else m.__name__
    m = m if isinstance(m,ModuleType) else sys.modules[m]
    
    # We need to have a reference to each of the old classes because the
    # reload that we are about to do will overwrite their names in the module.
    # here we create a dict with (key,value)=(class name, class reference):
    oldClasses = {}
    for key in dir(m):
        val = getattr(m,key)
        if isinstance(val,ClassType) or getattr(val,'__class__',False) is type:
            oldClasses[key] = val 
            
    # Ok, now we can do the reload.
    reload(m)  
    
    # Now that we've done the reload we can get a dict of all the 
    # currently decorated classes, of the same form as oldClasses
    taggedNewClasses = {}
    for key in dir(m):
        val = getattr(m,key)
        if getattr(val,'__module__',None) == name_m and \
                hasattr(val,'__CAN_RELOAD_METHODS__'):
            taggedNewClasses[key] = val
        
    for name_c,c in taggedNewClasses.iteritems():
        if name_c not in oldClasses:
            continue # if we cant find the old version of the class then there's nothing much we can do.

        # For the given class, we don't just want to update the methods in oldClasses[name_c]
        # we also need to follow the chain back through all previous versions of the class
        # ("all" meaning all versions that were tagged plus the version just before it was tagged)            
        c.__PREVIOUS_VERSION__ = oldClasses[name_c]   
        c_chain = [c.__PREVIOUS_VERSION__]
        while hasattr(c_chain[-1],'__PREVIOUS_VERSION__'):
            c_chain.append(c_chain[-1].__PREVIOUS_VERSION__)

        # Right, lets get a list of all "method"-like things for the new version of the class
        # These get stored in a dict as (key,value)=(name,reference)
        # Note that we collect the im_func of methods, and the whole of properties
        newMethods = {}
        newProperties = {}
        for attr_name in dir(c):
            attr = getattr(c,attr_name)
            if hasattr(attr,'im_func'):
                newMethods[attr_name] = attr.im_func
            elif isinstance(attr,property):
                newProperties[attr_name] = attr
                        
        # Ok, finally we can now iterate over all previous versions of the class
        # and override the old/non-existent attr with the new vesion of the 
        # methods/properties:
        for met_name,met in newMethods.iteritems():
            for c_old in c_chain:
                setattr(c_old,met_name,_enclose(met))
        for prop_name,prop in newProperties.iteritems():
            for c_old in c_chain:
                setattr(c_old,prop_name,prop)                                
        
        
        print "Updated %d methods and %d properties of %d versions of class %s in module %s." % \
                (len(newMethods),len(newProperties),len(c_chain),name_c,name_m)
        


def _enclose(met): #builds closure around iterator
    def wrapped(*args,**kwargs):
        return met(*args,**kwargs)
    return wrapped