modulemaskcoveragewhore.py

import os
import sys


class ModuleMask(object):
    """Hide a module, so it behaves as it's not installed.
    
    This is for the coverage whores out there.
    
    Examples usage
    

        >>> del(sys.modules["celery.serialization"])
        >>> with ModuleMask("cPickle") as modmask:
        >>>    from celery.serialization import pickle
        >>>    import pickle as orig_pickle
        >>>    self.assertTrue(pickle.dumps is orig_pickle.dumps)

    """
   
    def __init__(self, module_name):
        self.module_name = module_name
        self.modpath = None
        self.modpath_pos = None
        self.mod = None

    def __enter__(self):
        module_name = self.module_name
        mod = __import__(module_name, {}, {}, [''])
        modpath = os.path.dirname(mod.__file__)
        modpath_pos_found = False
        for modpath_pos, syspath in enumerate(sys.path):
            if syspath == modpath:
                modpath_pos_found = True
                break
        if not modpath_pos_found:
            raise Exception("Module path not in sys.path")
       
        self.mod = mod
        self.modpath = modpath
        self.modpath_pos = modpath_pos
        del(sys.path[modpath_pos])
        del(sys.modules[module_name])
        return self
    def __exit__(self, e_type, e_value, e_trace):
        sys.path.insert(self.modpath_pos, self.modpath)
        sys.modules[self.module_name] = self.mod
        if e_type:
            raise e_type(e_value)