virtualenv.py

import os
import sys

import click

from pyinfra import logger


def init_virtualenv():
    '''
    Add a virtualenv to sys.path so the user can import modules from it.
    This isn't perfect: it doesn't use the Python interpreter with which the
    virtualenv was built, and it ignores the --no-site-packages option. A
    warning will appear suggesting the user installs IPython in the
    virtualenv, but for many cases, it probably works well enough.

    Adapted from IPython's implementation:
    https://github.com/ipython/ipython/blob/master/IPython/core/interactiveshell.py
    '''

    if 'VIRTUAL_ENV' not in os.environ:
        # Not in a virtualenv
        return

    p = os.path.normcase(sys.executable)
    p_venv = os.path.normcase(os.environ['VIRTUAL_ENV'])

    # executable path should end like /bin/python or \\scripts\\python.exe
    p_exe_up2 = os.path.dirname(os.path.dirname(p))
    if p_exe_up2 and os.path.samefile(p_exe_up2, p_venv):
        # Our exe is inside the virtualenv, don't need to do anything.
        return

    # fallback venv detection:
    # stdlib venv may symlink sys.executable, so we can't use realpath.
    # but others can symlink *to* the venv Python, so we can't just use sys.executable.
    # So we just check every item in the symlink tree (generally <= 3)
    paths = [p]
    while os.path.islink(p):
        p = os.path.normcase(os.path.join(os.path.dirname(p), os.readlink(p)))
        paths.append(p)

    # In Cygwin paths like 'c:\...' and '\cygdrive\c\...' are possible
    if p_venv.startswith('\\cygdrive'):
        p_venv = p_venv[11:]
    elif len(p_venv) >= 2 and p_venv[1] == ':':
        p_venv = p_venv[2:]

    if any(p_venv in p for p in paths):
        # Running properly in the virtualenv, don't need to do anything
        return

    logger.warning((
        'Attempting to work in a virtualenv.\n'
        '    If you encounter problems, please install pyinfra inside the virtualenv.'
    ))
    click.echo(err=True)

    if sys.platform == 'win32':
        virtual_env = os.path.join(
            os.environ['VIRTUAL_ENV'],
            'Lib',
            'site-packages',
        )
    else:
        virtual_env = os.path.join(
            os.environ['VIRTUAL_ENV'],
            'lib',
            'python%d.%d' % sys.version_info[:2],
            'site-packages',
        )

    import site
    sys.path.insert(0, virtual_env)
    site.addsitedir(virtual_env)