recipe_modules/environment/api.py

# Copyright 2019 The Pigweed Authors
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.
"""Environment utility functions.

Usage:
api.environment.initialize(checkout_root=...)
with api.environment():
  ...
"""

import attr
from recipe_engine import recipe_api


@attr.s
class Package(object):
  name = attr.ib(type=bytes)
  version = attr.ib(type=bytes)


class EnvironmentApi(recipe_api.RecipeApi):
  """Environment utility functions."""

  def __init__(self, *args, **kwargs):
    super(EnvironmentApi, self).__init__(*args, **kwargs)
    self._cipd_dir = None
    self._path_prefixes = []
    self._ldpath_prefixes = []
    self._env = {}
    self._initialized = False

  def _process_ensure_path(self, ensure_path, ensure_file):
    """Add packages from ensure_path to ensure_file for later installation."""

    ensure_text = self.m.file.read_text(
        'read {}'.format(self.m.path.basename(ensure_path)), ensure_path)
    packages = []
    for line in ensure_text.splitlines():
      line = line.strip()
      if not line or line[0] in '$#':
        continue
      packages.append(Package(*line.split()))

    for package in packages:
      ensure_file.add_package(package.name, package.version)

  def _init_cipd(self, checkout_root):
    """Install CIPD packages."""

    with self.m.step.nest('setup cipd'):
      ensure_paths = [
          checkout_root.join('env_setup/cipd/pigweed.ensure'),
          # TODO(mohrr) handle multiple files in non-predetermined places
      ]

      ensure_file = self.m.cipd.EnsureFile()
      for ensure_path in ensure_paths:
        self._process_ensure_path(ensure_path, ensure_file)

      self._cipd_dir = self.m.path['start_dir'].join('cipd')
      self.m.cipd.max_threads = 0
      self.m.cipd.ensure(self._cipd_dir, ensure_file)

      self._path_prefixes.append(self._cipd_dir)
      self._path_prefixes.append(self._cipd_dir.join('bin'))
      self._ldpath_prefixes.append(self._cipd_dir)
      self._ldpath_prefixes.append(self._cipd_dir.join('lib'))
      self._env['PW_CIPD_INSTALL_DIR'] = self._cipd_dir
      self._env['PW_PIGWEED_CIPD_INSTALL_DIR'] = self._cipd_dir

  def _all_python_packages(self, checkout_root):
    """Return all folders with setup.py entries.

    Note: this function should not be called from within a virtualenv. It screws
    with how vpython works.

    Args:
      checkout_root(Path): root of source tree

    Returns:
      A list of package paths.
    """
    files = self.m.file.listdir('ls **/setup.py', checkout_root, recursive=True)
    matches = [
        self.m.path.dirname(x)
        for x in files
        if self.m.path.basename(x) == 'setup.py'
    ]
    with self.m.step.nest('packages') as step:
      step.logs['matches'] = [str(x) for x in matches]
    return matches

  def _init_python(self, checkout_root):
    """Initialize the Python environment. (Specifically, install 'pw'.)"""

    with self.m.step.nest('setup python'):
      packages = self._all_python_packages(checkout_root)

      with self.m.step.nest('setup virtualenv'), self():
        venv_dir = self.m.path['start_dir'].join('venv')
        venv_bin = venv_dir.join('bin')
        python = venv_bin.join('python3')
        self.m.step('create venv', ['python3', '-m', 'venv', venv_dir])
        self.m.step('upgrade pip',
                    [python, '-m', 'pip', 'install', '--upgrade', 'pip'])

        # Need to insert at beginning because venv python should trump cipd
        # python.
        self._path_prefixes.insert(0, venv_bin)
        self._env['VIRTUAL_ENV'] = venv_dir

      # Need to exit and reenter 'with self()' to include new context from
      # creating the virtualenv.

      with self.m.step.nest('install packages'), self():
        pip_install_prefix = (python, '-m', 'pip', 'install')
        pw_cmd = list(pip_install_prefix)
        for package in packages:
          pw_cmd.append('--editable={}'.format(package))
        self.m.step('pigweed tools', pw_cmd)

        requirements = checkout_root.join(
            'env_setup/virtualenv/requirements.txt')
        req_cmd = list(pip_install_prefix)
        req_cmd.extend(('-r', requirements))
        self.m.step('build requirements', req_cmd)

  def init(self, checkout_root):
    if not self._initialized:
      with self.m.step.nest('environment'):
        # Setting _initialized immediately because some setup steps need to
        # use the context of previous steps, and invoking self() is the
        # easiest way to do so.
        self._initialized = True
        self._env['PW_ROOT'] = checkout_root
        self._init_cipd(checkout_root)
        self._init_python(checkout_root)

  def __call__(self):
    assert self._initialized

    env_prefixes = {}
    if self._path_prefixes:
      env_prefixes['PATH'] = self._path_prefixes
    if self._ldpath_prefixes:
      env_prefixes['LD_LIBRARY_PATH'] = self._ldpath_prefixes

    return self.m.context(env_prefixes=env_prefixes, env=self._env)