core/homeassistant/util/package.py

71 lines
2.2 KiB
Python

"""Helpers to install PyPi packages."""
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, target):
return True
_LOGGER.info("Attempting install of %s", package)
args = [sys.executable, '-m', 'pip', 'install', '--quiet', package]
if upgrade:
args.append('--upgrade')
if target:
args += ['--target', os.path.abspath(target)]
if constraints is not None:
args += ['--constraint', constraints]
process = Popen(args, stdin=PIPE, stdout=PIPE, stderr=PIPE)
_, 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, lib_dir: 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)
# Check packages from lib dir
if lib_dir is not None:
if any(dist in req for dist in
pkg_resources.find_distributions(lib_dir)):
return True
# Check packages from global + virtual environment
# pylint: disable=not-an-iterable
return any(dist in req for dist in pkg_resources.working_set)