recipe_modules/environment/api.py

# Copyright 2020 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 contextlib

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, props, *args, **kwargs):
        super(EnvironmentApi, self).__init__(*args, **kwargs)
        self._cargo_package_files = props.cargo_package_files
        self._cipd_package_files = props.cipd_package_files
        self._virtualenv_requirements = props.virtualenv_requirements
        self._virtualenv_setup_py_roots = props.virtualenv_setup_py_roots
        self._root_variable_name = str(props.root_variable_name)
        self._relative_pigweed_root = str(props.relative_pigweed_root)
        self._cipd_dir = None
        self._path_prefixes = []
        self._ldpath_prefixes = []
        self._env = {}
        self._initialized = False
        self._cipd_installation_dirs = []

    def _convert_paths(self, root, pathlist):
        results = []
        for path in pathlist:
            parts = [x for x in path.split('/') if x != '.']
            results.append(root.join(*parts))
        return results

    def _init_platform(self):
        if self.m.platform.is_mac:
            with self.m.step.nest('setup platform'):
                with self.m.macos_sdk():
                    pass

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

        with self.m.step.nest('setup cipd'):
            self._cipd_dir = self.m.path['start_dir'].join('cipd')
            self._env['PW_CIPD_INSTALL_DIR'] = self._cipd_dir

            for json_path in self._convert_paths(
                checkout_root, self._cipd_package_files
            ):
                name = self.m.path.splitext(self.m.path.basename(json_path))[0]
                with self.m.step.nest(name):
                    install_dir = self._cipd_dir.join(name)
                    self._env[
                        'PW_{}_CIPD_INSTALL_DIR'.format(name.upper())
                    ] = install_dir

                    packages = self.m.file.read_json(
                        'read {}'.format(json_path), json_path
                    )
                    if not packages:
                        continue

                    ensure_file = self.m.cipd.EnsureFile()
                    for pkg in packages:
                        # JSON files are read as unicode, need to encode to get
                        # str.
                        ensure_file.add_package(
                            pkg['path'].encode(),
                            ' '.join(x.encode() for x in pkg['tags']),
                            subdir=pkg.get('subdir'),
                        )
                    self.m.cipd.ensure(install_dir, ensure_file)

                    for path in ('', 'bin', 'mingw64/bin'):
                        path = install_dir.join(*path.split('/'))
                        if self.m.path.exists(path):
                            self._path_prefixes.append(path)

                    for path in ('', 'lib'):
                        path = install_dir.join(*path.split('/'))
                        if self.m.path.exists(path):
                            self._ldpath_prefixes.append(path)

                self._cipd_installation_dirs.append(install_dir)

    def _find_python_packages(self, roots):
        """Return all folders with setup.py entries.

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

        Args:
            roots(Collection[Path]): locations to search for setup.py files

        Returns:
            A list of package paths.
        """
        matches = []
        with self.m.step.nest('find_python_packages'):
            for root in roots:
                files = self.m.file.listdir(str(root), root, recursive=True)
                matches.extend(
                    self.m.path.dirname(x)
                    for x in files
                    if self.m.path.basename(x) == 'setup.py'
                )

            original_matches = matches[:]
            for match in original_matches:
                # Remove all matches that have another match as a strict prefix.
                matches = [
                    x for x in matches if x == match or not x.startswith(match)
                ]

            with self.m.step.nest('packages') as step:
                step.logs['raw matches'] = [str(x) for x in original_matches]
                step.logs['filtered 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._find_python_packages(
                self._convert_paths(
                    checkout_root, self._virtualenv_setup_py_roots
                )
            )

            with self.m.step.nest('setup virtualenv'), self():
                venv_dir = self.m.path['start_dir'].join('venv')
                cipd_python = 'python3'

                # Work around weird bug where venv creation requires the name of
                # the executable to be python.exe and not python3.exe on
                # Windows.
                if self.m.platform.is_win:
                    for install_dir in self._cipd_installation_dirs:
                        this_python = install_dir.join('bin', 'python3.exe')
                        if self.m.path.exists(this_python):
                            cipd_python = this_python
                            new_cipd_python = install_dir.join(
                                'bin', 'python.exe'
                            )
                            self.m.file.copy(
                                'cp python3.exe python.exe',
                                cipd_python,
                                new_cipd_python,
                            )
                            cipd_python = new_cipd_python

                self.m.step(
                    'create venv', [cipd_python, '-m', 'venv', venv_dir]
                )

                venv_bin = venv_dir.join('bin')
                if not self.m.path.exists(venv_bin):
                    alt_venv_bin = venv_dir.join('Scripts')
                    # TODO(mohrr) figure out how to cover this.
                    if self.m.path.exists(alt_venv_bin):  # pragma: no cover
                        venv_bin = alt_venv_bin

                python = venv_bin.join('python')
                self.m.step(
                    'upgrade pip',
                    [python, '-m', 'pip', 'install', '--upgrade', 'pip'],
                )

                self._path_prefixes.append(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)

                for req in self._convert_paths(
                    checkout_root, self._virtualenv_requirements
                ):
                    req_cmd = list(pip_install_prefix)
                    req_cmd.extend(('-r', req))
                    self.m.step('build requirements', req_cmd)

    def _init_cargo(self, checkout_root):
        if not self._cargo_package_files:
            return

        prefix = self.m.path['start_dir'].join('cargo')
        self._path_prefixes.append(prefix.join('bin'))
        cache = self.m.path['cache'].join('cargo')
        self._env['CARGO_TARGET_DIR'] = cache

        # On desktops this is the cue to include cargo in setup, but other
        # things assume this variable can tell them if we're set up with cargo,
        # so set it here for those checks.
        self._env['PW_CARGO_SETUP'] = '1'

        with self.m.step.nest('setup cargo'):
            for packages_txt in self._convert_paths(
                checkout_root, self._cargo_package_files
            ):
                raw = self.m.file.read_text(
                    'read {}'.format(packages_txt), packages_txt
                )

                for line in raw.splitlines():
                    line = line.strip()
                    if not line or line.startswith('#'):
                        continue

                    package, version = line.split()
                    cmd = [
                        'cargo',
                        'install',
                        '--force',
                        '--root',
                        prefix,
                        '--version',
                        version,
                        package,
                    ]

                    # Deliberately narrow scope--read_text() above fails if run
                    # within 'with self()' because Python virtual environment
                    # includes Python 3 as 'python' in PATH.
                    with self():
                        self.m.step('install {}'.format(package), cmd)

    def _init_misc(self):
        self._env['GOCACHE'] = self.m.path['cache'].join('go')

    def init(self, checkout_root):
        pigweed_root = checkout_root
        if self._relative_pigweed_root not in (None, '', '.'):
            pigweed_root = checkout_root.join(
                *self._relative_pigweed_root.split('/')
            )

        if self._root_variable_name:
            self._env[self._root_variable_name] = 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'] = pigweed_root

                self._init_platform()
                self._init_cipd(checkout_root)
                self._init_python(checkout_root)
                self._init_cargo(checkout_root)
                self._init_misc()

    @contextlib.contextmanager
    def __call__(self):
        assert self._initialized

        env_prefixes = {}
        # Using reversed() because things that are added later in environment
        # setup need to override things that came earlier.
        if self._path_prefixes:
            env_prefixes['PATH'] = reversed(self._path_prefixes)
        if self._ldpath_prefixes:
            env_prefixes['LD_LIBRARY_PATH'] = reversed(self._ldpath_prefixes)

        with self.m.context(env_prefixes=env_prefixes, env=self._env):
            with self.m.macos_sdk():
                yield self

    def __getattr__(self, name):
        if name not in self._env:
            raise AttributeError(name)

        return self._env.get(name)