[three]Bean

Switching virtualenvs with a python context manager

Jun 06, 2011 | categories: pip, python, virtualenv, pypi View Comments

EDIT: I released this. You can find it on pypi.

---

Got it! Spent the last day working on a control script for moksha to replace my mistake of choosing fabric.

With this little nugget, you can do cool in-python context switching of virtualenvs with VirtualenvContext like this:

#!/usr/bin/python

from blogcopypaste import VirtualenvContext

try:
    import kitchen
except ImportError as e:
    print "kitchen is definitely not installed in system-python"

with VirtaulenvContext("my-venv"):
    import kitchen
    print "But it *is* installed in my virtualenv"

try:
    import kitchen
except ImportError as e:
    print "But once I exit that block, I lose my powers again..."

kitchen could be any non-standard-library python package you choose. (Although kitchen itself is pretty cool).

I learned a ton about ihooks and the python __import__ built-in... PEP 302 was an eye-opener.

Here's the code that makes that fancy VirtualenvContext happen:

""" Virtualenv context management! """

import os
import sys
import ihooks
import warnings
import imp

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):
    """ Overridden 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):
        workon = os.getenv("WORKON_HOME", None)
        venv_location = "/".join([
            workon, self.venv, 'lib/python2.7/site-packages'])
        full = lambda i : "/".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_name):
        self.venv = venv_name
        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
        workon = os.getenv("WORKON_HOME", None)
        sys.executable = "/".join([workon, 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()

Fun fact: you can combine this with the install_distributions function in my previous post to do:

with VirtualenvContext('some-environment'):
    install_distributions(['Markdown'])

View Comments