Source

virtualenvmanager / virtualenvmanager.py

Full commit
"""
    envsetup
    ~~~~~~~~

    a simple lib for creating virtualenvs and setting up versions of libraries

"""
import os
import subprocess
import commands
try:
    from subprocess import CalledProcessError
    from subprocess import check_call
except:
    class CalledProcessError(Exception):
        pass
    def check_call(args, **kw):
        retcode = subprocess.call(args, **kw)
        if retcode != 0:
            raise CalledProcessError(retcode, *args)


class VirtualEnvBase(object):

    def create(self, sitepackages=False, python_executable=None):
        args = ['virtualenv', self.path]
        if not sitepackages:
            args.append('--no-site-packages')
        if python_executable:
            args.append('--python=%s'%python_executable)
        self._check_call(args)

    def _cmd(self, name):
        return os.path.join(self.path, 'bin', name)

    @property
    def valid(self):
        return self._has_cmd('python')

    @property
    def has_pip(self):
        return self._has_cmd('pip')

    def pcall(self, cmd, *args, **kw):
        assert self.valid
        arg = [self._cmd(cmd)]
        arg.extend(args)
        return self._subprocess_result(arg, **kw)

    def setup(self, path, action):
        #XXX: extend?
        from os.path import join
        return self.pcall('python', join(path, 'setup.py'), action, cwd=path)

    def pip_install(self, *packages):
        if not self.has_pip:
            self.easy_install('pip')
        return self.pcall('pip', 'install', *packages)

    def easy_install(self, *packages, **kw):
        args = []
        if 'index' in kw:
            index = kw['index']
            if isinstance(index, (list, tuple)):
                for i in index:
                    args.extend(['-i', i])
            else:
                args.extend(['-i', index])

        args.extend(packages)
        return self.pcall('easy_install', *args)

    def makegateway(self):
        '''
        pseudo-fix sys.path in the venv
        '''
        import execnet
        gw = execnet.makegateway(self.spec)
        channel = gw.remote_exec("""
            import sys
            sys.path.insert(0, sys.path[-1])
            channel.send(True)
            """)
        assert channel.receive(), """failed to activate virtualenv"""
        return gw

    def has_package(self, name):
        gw = self.makegateway()
        path = gw.remote_exec('''
            import sys, os
            import %s as mod
            file = os.path.abspath(mod.__file__)
            channel.send(file)
        '''%name).receive()

        return path.startswith(self.path)

    def has_distribution(self, dist):
        gw = self.makegateway()
        channel = gw.remote_exec('''
            dist = channel.receive()
            import pkg_resources
            try:
                pkg_resources.get_distribution(dist)
                channel.send(True)
            except pkg_resources.DistributionNotFound:
                channel.send(False)
        ''')
        channel.send(dist)
        return channel.receive()

    def ensure_distribution(self, dist):
        if self.has_distribution(dist):
            return True
        else:
            #XXX: when can we default to pip_install?
            return not self.easy_install(dist)

class VirtualEnv(VirtualEnvBase):
    def __init__(self, path, ):
        #XXX: supply the python executable
        self.path = path

    def __repr__(self):
        return "<VirtualEnv at %r>" %(self.path)

    def ensure(self):
        if not self.valid:
            self.create()

    def _check_call(self, args):
        check_call(args)

    def _subprocess_result(self, args, **kw):
        return subprocess.call(args, **kw)

    def _has_cmd(self, cmd):
        return os.path.exists(self._cmd(cmd))

    @property
    def python_version(self):
        #XXX: use pylib?
        cmd = '%s -u -c "import sys;print(sys.version)"'
        return commands.getoutput(cmd%self._cmd('python')).strip()

    @property
    def spec(self):
        return 'popen//python=%s//chdir=%s'%(
                    self._cmd('python'),
                    self.path,
                )


class RemoteVirtualenv(VirtualEnvBase):
    '''
    Manages a virtualenv within the working directory of a execnet gateway
    '''
    def __init__(self, gw, name):
        self.gw = gw
        self.name = name

    @property
    def path(self):
        base = self.gw.remote_exec('''
            import os
            channel.send(os.getcwd())
            ''').receive()
        return os.path.join(base, self.name)

    @property
    def python_version(self):
        return self.makegateway().remote_exec('''
            import sys
            channel.send(sys.version)
            ''').receive()

    def _subprocess_result(self, args, **kw):
        channel = self.gw.remote_exec('''
            import subprocess
            args, kw = channel.receive()
            retcode = subprocess.call(args, **kw)
            channel.send(retcode)
            ''')
        channel.send((args, kw))
        return channel.receive()

    def _check_call(self, args):
        retcode = self._subprocess_result(args)
        if retcode != 0:
            raise CalledProcessError(retcode, *args)

    def _has_cmd(self, cmd):
        cmd = self._cmd(cmd)
        channel = self.gw.remote_exec('''
            import os
            channel.send(os.path.exists(channel.receive()))
            ''')
        channel.send(cmd)
        return channel.receive()
    
    @property
    def spec(self):
        from execnet.xspec import XSpec
        meta = XSpec(str(self.gw.spec))
        meta.chdir=self.path
        meta.python=self._cmd('python')
        if meta.ssh:
            #XXX: untested
            base = 'ssh=%(ssh)//python=%(python)s//chdir=%(chdir)s'
        elif meta.popen:
            base = 'popen//python=%(python)s//chdir=%(chdir)s'
        else:
            raise RuntimeError('wrong xspec', str(meta))
        return base % vars(meta)