"""Helpers to install PyPi packages.""" import asyncio import logging import os import sys import threading from subprocess import Popen, PIPE from urllib.parse import urlparse from typing import Optional import pkg_resources _LOGGER = logging.getLogger(__name__) INSTALL_LOCK = threading.Lock() def install_package(package: str, upgrade: bool=True, target: Optional[str]=None, constraints: Optional[str]=None) -> bool: """Install a package on PyPi. Accepts pip compatible package strings. Return boolean if install successful. """ # Not using 'import pip; pip.main([])' because it breaks the logger with INSTALL_LOCK: if check_package_exists(package): return True _LOGGER.info('Attempting install of %s', package) env = os.environ.copy() args = [sys.executable, '-m', 'pip', 'install', '--quiet', package] if upgrade: args.append('--upgrade') if constraints is not None: args += ['--constraint', constraints] if target: assert not is_virtual_env() # This only works if not running in venv args += ['--user'] env['PYTHONUSERBASE'] = os.path.abspath(target) if sys.platform != 'win32': # Workaround for incompatible prefix setting # See http://stackoverflow.com/a/4495175 args += ['--prefix='] process = Popen(args, stdin=PIPE, stdout=PIPE, stderr=PIPE, env=env) _, stderr = process.communicate() if process.returncode != 0: _LOGGER.error("Unable to install package %s: %s", package, stderr.decode('utf-8').lstrip().strip()) return False return True def check_package_exists(package: str) -> bool: """Check if a package is installed globally or in lib_dir. Returns True when the requirement is met. Returns False when the package is not installed or doesn't meet req. """ try: req = pkg_resources.Requirement.parse(package) except ValueError: # This is a zip file req = pkg_resources.Requirement.parse(urlparse(package).fragment) env = pkg_resources.Environment() return any(dist in req for dist in env[req.project_name]) def is_virtual_env() -> bool: """Return true if environment is a virtual environment.""" return hasattr(sys, 'real_prefix') def _get_user_site(deps_dir: str) -> tuple: """Get arguments and environment for subprocess used in get_user_site.""" env = os.environ.copy() env['PYTHONUSERBASE'] = os.path.abspath(deps_dir) args = [sys.executable, '-m', 'site', '--user-site'] return args, env def get_user_site(deps_dir: str) -> str: """Return user local library path.""" args, env = _get_user_site(deps_dir) process = Popen(args, stdin=PIPE, stdout=PIPE, stderr=PIPE, env=env) stdout, _ = process.communicate() lib_dir = stdout.decode().strip() return lib_dir @asyncio.coroutine def async_get_user_site(deps_dir: str, loop: asyncio.AbstractEventLoop) -> str: """Return user local library path. This function is a coroutine. """ args, env = _get_user_site(deps_dir) process = yield from asyncio.create_subprocess_exec( *args, loop=loop, stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.DEVNULL, env=env) stdout, _ = yield from process.communicate() lib_dir = stdout.decode().strip() return lib_dir