[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

Installing from pip inside python or, a simple pip API

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

So, it's a long story, but my bid to convert moksha's development tools to fabric was a terrible idea: mandatory sshd on dev boxes, passwords in the clear when chaining commands between systems, simple failure to work. It was the wrong tool for the job. Now I have my work cut out for me to replace it with something worthwhile.

I have a need to invoke pip from inside python and to do it within the current virtualenv. I googled and found this somewhat old thread on the virtualenv list which indicates that there's not much help out there.

Here's my stab at it.

import pip.commands.install

def install_distributions(distributions):
    command = pip.commands.install.InstallCommand()
    opts, args = command.parser.parse_args()
    # TBD, why do we have to run the next part here twice before actual install
    requirement_set = command.run(opts, distributions)
    requirement_set = command.run(opts, distributions)
    requirement_set.install(opts)

And a test to see if it works.

#!/usr/bin/env python

from blogmodule import install_distributions

def run_test():
    try:
        import markdown
        print "Markdown is installed!  Aborting."
        assert(False)
    except ImportError as e:
        print "Markdown isn't yet installed.  That's good."

    install_distributions(["Markdown"])

    try:
        import markdown
        print "Markdown is now installed.  That's good!"
    except ImportError as e:
        print "Markdown never got installed.  That's bad."
        assert(False)

if __name__ == '__main__':
    run_test()

Looking forward to doing something cool with this and python context managers to manage virtualenvs.

View Comments