vc.py

vc.py

'''
vc.py is a slight modification of virtualenvcontext, and it offers the possibility
to use the context manager without virtualenvwrapper.

This is useful while testing things, because it allows you to works with ad hoc
virtualenvs.
'''

## Modified from virtualenvcontext (C) 2011 Ralph Bean
## http://pypi.python.org/pypi/virtualenvcontext/


import os
import sys
import imp
import ihooks
import warnings


def _silent_load_source(name, filename, file=None):
    '''Helper function.  Overrides a import hook.  Suppresses warnings.'''
    with warnings.catch_warnings():
        warnings.simplefilter('ignore')
        return imp.load_source(name, filename, file)


class VenvModuleLoader(ihooks.ModuleLoader):
    '''Checks for a virtualenv first and remembers imports.'''

    remembered = []

    def __init__(self, venv, verbose=0):
        self.venv = venv
        ihooks.ModuleLoader.__init__(self, verbose=verbose)
        self.hooks.load_source = _silent_load_source

    def default_path(self):
        venv_location = os.path.join([self.venv, 'lib/python2.7/site-packages'])
        full = lambda i: os.path.join([venv_location, i])
        venv_path = [venv_location] + [
            full(item) for item in os.listdir(venv_location)
            if os.path.isdir(full(item))] + sys.path
        return venv_path + sys.path

    def load_module(self, name, stuff):
        ''' Overloaded just to remember what we load '''
        self.remembered.append(name)
        return ihooks.ModuleLoader.load_module(self, name, stuff)


class VirtualenvContext(object):
    '''Context manager for entering a virtualenv'''

    def __init__(self, venv_path):
        self.venv = venv_path
        self.loader = VenvModuleLoader(venv=self.venv)
        self.importer = ihooks.ModuleImporter(loader=self.loader)

    def __enter__(self):
        # Install our custom importer
        self.importer.install()

        # Pretend like our exectuable is really somewhere else
        self.old_exe = sys.executable
        sys.executable = os.path.join([self.venv, 'bin/python'])

    def __exit__(self, exc_type, exc_value, traceback):
        # Uninstall our custom importer
        self.importer.uninstall()

        # Reset our executable
        sys.exectuable = self.old_exe

        # Unload anything loaded while inside the context
        for name in self.importer.loader.remembered:
            if not name in sys.modules:
                continue
            del sys.modules[name]
        self.importer.loader.remembered = []
        sys.path_importer_cache.clear()