ha-core/homeassistant/requirements.py

"""Module to handle installing requirements."""
import asyncio
from functools import partial
import logging
import os
import sys
from typing import Any, Dict, List, Optional
from urllib.parse import urlparse

import pkg_resources

import homeassistant.util.package as pkg_util
from homeassistant.core import HomeAssistant

DATA_PIP_LOCK = 'pip_lock'
DATA_PKG_CACHE = 'pkg_cache'
CONSTRAINT_FILE = 'package_constraints.txt'
_LOGGER = logging.getLogger(__name__)


async def async_process_requirements(hass: HomeAssistant, name: str,
                                     requirements: List[str]) -> bool:
    """Install the requirements for a component or platform.

    This method is a coroutine.
    """
    pip_lock = hass.data.get(DATA_PIP_LOCK)
    if pip_lock is None:
        pip_lock = hass.data[DATA_PIP_LOCK] = asyncio.Lock(loop=hass.loop)

    pkg_cache = hass.data.get(DATA_PKG_CACHE)
    if pkg_cache is None:
        pkg_cache = hass.data[DATA_PKG_CACHE] = PackageLoadable(hass)

    pip_install = partial(pkg_util.install_package,
                          **pip_kwargs(hass.config.config_dir))

    async with pip_lock:
        for req in requirements:
            if await pkg_cache.loadable(req):
                continue

            ret = await hass.async_add_executor_job(pip_install, req)

            if not ret:
                _LOGGER.error("Not initializing %s because could not install "
                              "requirement %s", name, req)
                return False

    return True


def pip_kwargs(config_dir: Optional[str]) -> Dict[str, Any]:
    """Return keyword arguments for PIP install."""
    kwargs = {
        'constraints': os.path.join(os.path.dirname(__file__), CONSTRAINT_FILE)
    }
    if not (config_dir is None or pkg_util.is_virtual_env()):
        kwargs['target'] = os.path.join(config_dir, 'deps')
    return kwargs


class PackageLoadable:
    """Class to check if a package is loadable, with built-in cache."""

    def __init__(self, hass: HomeAssistant) -> None:
        """Initialize the PackageLoadable class."""
        self.dist_cache = {}  # type: Dict[str, pkg_resources.Distribution]
        self.hass = hass

    async def loadable(self, package: str) -> bool:
        """Check if a package is what will be loaded when we import it.

        Returns True when the requirement is met.
        Returns False when the package is not installed or doesn't meet req.
        """
        dist_cache = self.dist_cache

        try:
            req = pkg_resources.Requirement.parse(package)
        except ValueError:
            # This is a zip file. We no longer use this in Home Assistant,
            # leaving it in for custom components.
            req = pkg_resources.Requirement.parse(urlparse(package).fragment)

        req_proj_name = req.project_name.lower()
        dist = dist_cache.get(req_proj_name)

        if dist is not None:
            return dist in req

        for path in sys.path:
            # We read the whole mount point as we're already here
            # Caching it on first call makes subsequent calls a lot faster.
            await self.hass.async_add_executor_job(self._fill_cache, path)

            dist = dist_cache.get(req_proj_name)
            if dist is not None:
                return dist in req

        return False

    def _fill_cache(self, path: str) -> None:
        """Add packages from a path to the cache."""
        dist_cache = self.dist_cache
        for dist in pkg_resources.find_distributions(path):
            dist_cache.setdefault(dist.project_name.lower(), dist)