virtualenvcontext.py

import os
import shlex
import subprocess
from subprocess import Popen
from pathlib import Path


class VirtualEnv(object):
    """
    Create a VirtualEnv context manager.
    It changes the environment variables to that of the virtualenv through bin/activate on __enter__
    and returns it back to the previous state on __exit__.
    
    Example use:
    my_virtual_env = VirtualEnv("/path/to/my/virtualenv")
    with VirtualEnv(my_virtual_env) as env:
        cmd = ['script_or_executable_in_venv', '--arg1', '--arg2']
        p = subprocess.call(cmd) #this calls the subprocess as if running in the virtualenv
    
    """
    def __init__(self, path_to_virtualenv):
        path_to_virtualenv = Path(path_to_virtualenv).absolute()
        
        if not path_to_virtualenv.exists():
            raise NotADirectoryError(path_to_virtualenv)
        
        self.path_to_virtualenv = path_to_virtualenv
        
        self.venv_bin_activate = self.path_to_virtualenv / 'bin' / 'activate'
        self.python = self.path_to_virtualenv / 'bin' / 'python'
        
        if not self.venv_bin_activate.exists():
            raise FileNotFoundError(self.venv_bin_activate)
        
        if not self.python.exists():
            raise FileNotFoundError(self.python)
        
        cmd = '/bin/bash -c "source {} && python -c \'import os; print(dict(os.environ))\'"'
        self.cmd = cmd.format(str(self.venv_bin_activate))
        
        p = Popen(shlex.split(self.cmd), stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, shell=False)
        
        self.p = p
        
        self.env = eval(str(p.stdout.read(), 'utf-8'))
        
        self.current_env = os.environ.copy()
        
        self.p.terminate()
        
    def __getitem__(self, item):
        return self.env[item]
        
    def __enter__(self):
        os.environ.clear()
        os.environ.update(self.env)
        return self.env
    
    def __exit__(self ,type, value, traceback):
        os.environ.clear()
        os.environ.update(self.current_env)
        return self.current_env