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 pprint

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_gn_targets = props.virtualenv_gn_targets
        self._root_variable_name = str(props.root_variable_name)
        self._relative_pigweed_root = str(props.relative_pigweed_root)
        self._config_file = str(props.config_file)
        self._cipd_dir = None
        self._prefixes = {}
        self._suffixes = {}
        self._env = {}
        self._initialized = False
        self._cipd_installation_dirs = []

    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_misc_vars(self):
        self._env['PW_ENVSETUP_DISABLE_SPINNER'] = '1'

        if self.m.led.launched_by_led:
            # Not using self.m.buildbucket_util.id because having relatively
            # consistent length is important to some downstream projects when
            # pulling this into their builds.
            self._env['BUILDBUCKET_ID'] = 'test:{}'.format(
                self.m.swarming.task_id
            )
            self._env['BUILD_NUMBER'] = str(self.m.buildbucket.build.number)

        else:
            self._env['BUILDBUCKET_ID'] = str(self.m.buildbucket.build.id)
            self._env['BUILD_NUMBER'] = str(self.m.buildbucket.build.number)

        self._env['BUILDBUCKET_NAME'] = ':'.join(
            (
                self.m.buildbucket.build.builder.project,
                self.m.buildbucket.build.builder.bucket,
                self.m.buildbucket.build.builder.builder,
            )
        )

        if self._env['BUILDBUCKET_NAME'] == '::':
            self._env['BUILDBUCKET_NAME'] = 'project:bucket:builder'

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

    def _init_pigweed(self, checkout_root, top_pres):
        """Run pw_env_setup."""

        if (
            not self._cargo_package_files
            and not self._cipd_package_files
            and not self._virtualenv_requirements
            and not self._virtualenv_gn_targets
            and not self._config_file
        ):
            return

        def path(relative_path):
            parts = [
                x for x in relative_path.split('/') if x not in ('.', u'.')
            ]
            if parts:
                return checkout_root.join(*parts)
            else:
                return checkout_root  # pragma: no cover

        pw_root = path(self._relative_pigweed_root)

        env_dir = self.m.path['start_dir'].join('environment')
        json_file = env_dir.join('vars.json')
        shell_file = env_dir.join('setup.sh')
        venv_dir = env_dir.join('venv')

        self.m.file.ensure_directory(
            'mkdir {}'.format(self.m.path.basename(env_dir)), env_dir,
        )

        self.m.file.ensure_directory(
            'mkdir {}'.format(self.m.path.basename(venv_dir)), venv_dir,
        )

        cmd = [
            'python',
            pw_root.join(
                'pw_env_setup', 'py', 'pw_env_setup', 'env_setup.py'
            ),
            '--pw-root',
            pw_root,
            '--install-dir',
            env_dir,
            '--json-file',
            json_file,
            '--shell-file',
            shell_file,
            '--virtualenv-root',
            venv_dir,
        ]

        if self._config_file:
            cmd.extend(('--config-file', path(self._config_file)))

            top_pres.logs['vars.json'] = pprint.pformat(
                self.m.file.read_json('read config', path(self._config_file))
            )

        with self.m.step.nest('cipd package files') as pres:
            for pkg_file in self._cipd_package_files:
                cmd.extend(('--cipd-package-file', path(pkg_file)))
                pres.logs[pkg_file] = self.m.file.read_text(
                    pkg_file, path(pkg_file)
                )

        if self._cargo_package_files:
            cmd.append('--enable-cargo')
            with self.m.step.nest('cargo package files') as pres:
                for pkg_file in self._cargo_package_files:
                    cmd.extend(('--cargo-package-file', path(pkg_file)))
                    pres.logs[pkg_file] = self.m.file.read_text(
                        pkg_file, path(pkg_file)
                    )
                self._env['CARGO_TARGET_DIR'] = self.m.path['cache'].join(
                    'cargo'
                )

        with self.m.step.nest('virtualenv requirements') as pres:
            for req_file in self._virtualenv_requirements:
                cmd.extend(('--virtualenv-requirements', path(req_file)))
                pres.logs[req_file] = self.m.file.read_text(
                    req_file, path(req_file)
                )

        with self.m.step.nest('virtualenv gn targets') as pres:
            summary = []
            for target in self._virtualenv_gn_targets:
                parts = [x for x in target.relative_path.split('/')]
                path = checkout_root.join(*[x for x in parts if x != '.'])
                encoded = '{}#{}'.format(path, target.target)
                cmd.extend(('--virtualenv-gn-target', encoded))
                summary.append(encoded)
            pres.step_summary_text = '\n'.join(summary)

        with self.m.step.defer_results():
            with self.m.step.nest('run pw_env_setup'):
                with self():
                    self.m.step('pw_env_setup', cmd)

                env_files = self.m.file.listdir(
                    'ls {}'.format(self.m.path.basename(env_dir)), env_dir,
                ).get_result()

                venv_files = self.m.file.listdir(
                    'ls {}'.format(self.m.path.basename(venv_dir)),
                    venv_dir,
                ).get_result()

                for entry in env_files + venv_files:
                    ext = self.m.path.splitext(entry)[1]
                    if ext not in ('.cfg', '.log', '.sh', '.bat', '.json'):
                        continue

                    self.m.file.read_text(
                        'read {}'.format(self.m.path.basename(entry)),
                        entry,
                    )

        json_data = self.m.file.read_json(
            'read json file',
            json_file,
            test_data={
                'set': {'VIRTUAL_ENV': '/environment/virtualenv'},
                'modify': {
                    'PATH': {'append': ['/environment/bin']},
                    'LD_LIBRARY_PATH': {'prepend': ['/environment/lib']},
                },
            },
        )

        top_pres.logs['vars.json'] = pprint.pformat(json_data)

        for var, value in json_data['set'].iteritems():
            self._env[var] = value

        for var, actions in json_data['modify'].iteritems():
            for value in actions.get('prepend', ()):
                self._prefixes.setdefault(var, [])
                self._prefixes[var].append(value)
            for value in actions.get('append', ()):
                self._suffixes.setdefault(var, [])
                self._suffixes[var].append(value)

    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') as pres:
                # 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._env['PW_PROJECT_ROOT'] = checkout_root

                self._init_platform()
                self._init_misc_vars()
                self._init_pigweed(checkout_root, pres)

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

        # Using reversed() because things that are added later in environment
        # setup need to override things that came earlier.
        with self.m.context(
            env_prefixes={
                k: reversed(v) for k, v in self._prefixes.iteritems()
            },
            env_suffixes=self._suffixes,
            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)