1#!/usr/bin/env python
2
3# Copyright 2020 The Pigweed Authors
4#
5# Licensed under the Apache License, Version 2.0 (the "License"); you may not
6# use this file except in compliance with the License. You may obtain a copy of
7# the License at
8#
9#     https://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
13# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
14# License for the specific language governing permissions and limitations under
15# the License.
16"""Environment setup script for Pigweed.
17
18This script installs everything and writes out a file for the user's shell
19to source.
20
21For now, this is valid Python 2 and Python 3. Once we switch to running this
22with PyOxidizer it can be upgraded to recent Python 3.
23"""
24
25from __future__ import print_function
26
27import argparse
28import copy
29import glob
30import inspect
31import json
32import os
33import shutil
34import subprocess
35import sys
36
37# TODO(pwbug/67): Remove import hacks once the oxidized prebuilt binaries are
38# proven stable for first-time bootstrapping. For now, continue to support
39# running directly from source without assuming a functioning Python
40# environment when running for the first time.
41
42# If we're running oxidized, filesystem-centric import hacks won't work. In that
43# case, jump straight to the imports and assume oxidation brought in the deps.
44if not getattr(sys, 'oxidized', False):
45    old_sys_path = copy.deepcopy(sys.path)
46    filename = None
47    if hasattr(sys.modules[__name__], '__file__'):
48        filename = __file__
49    else:
50        # Try introspection in environments where __file__ is not populated.
51        frame = inspect.currentframe()
52        if frame is not None:
53            filename = inspect.getfile(frame)
54    # If none of our strategies worked, we're in a strange runtime environment.
55    # The imports are almost certainly going to fail.
56    if filename is None:
57        raise RuntimeError(
58            'Unable to locate pw_env_setup module; cannot continue.\n'
59            '\n'
60            'Try updating to one of the standard Python implemetations:\n'
61            '  https://www.python.org/downloads/')
62    sys.path = [
63        os.path.abspath(os.path.join(filename, os.path.pardir, os.path.pardir))
64    ]
65    import pw_env_setup  # pylint: disable=unused-import
66    sys.path = old_sys_path
67
68# pylint: disable=wrong-import-position
69from pw_env_setup.cipd_setup import update as cipd_update
70from pw_env_setup.cipd_setup import wrapper as cipd_wrapper
71from pw_env_setup.colors import Color, enable_colors
72from pw_env_setup import cargo_setup
73from pw_env_setup import environment
74from pw_env_setup import spinner
75from pw_env_setup import virtualenv_setup
76from pw_env_setup import windows_env_start
77
78
79# TODO(pwbug/67, pwbug/68) switch to shutil.which().
80def _which(executable,
81           pathsep=os.pathsep,
82           use_pathext=None,
83           case_sensitive=None):
84    if use_pathext is None:
85        use_pathext = (os.name == 'nt')
86    if case_sensitive is None:
87        case_sensitive = (os.name != 'nt' and sys.platform != 'darwin')
88
89    if not case_sensitive:
90        executable = executable.lower()
91
92    exts = None
93    if use_pathext:
94        exts = frozenset(os.environ['PATHEXT'].split(pathsep))
95        if not case_sensitive:
96            exts = frozenset(x.lower() for x in exts)
97        if not exts:
98            raise ValueError('empty PATHEXT')
99
100    paths = os.environ['PATH'].split(pathsep)
101    for path in paths:
102        try:
103            entries = frozenset(os.listdir(path))
104            if not case_sensitive:
105                entries = frozenset(x.lower() for x in entries)
106        except OSError:
107            continue
108
109        if exts:
110            for ext in exts:
111                if executable + ext in entries:
112                    return os.path.join(path, executable + ext)
113        else:
114            if executable in entries:
115                return os.path.join(path, executable)
116
117    return None
118
119
120class _Result:
121    class Status:
122        DONE = 'done'
123        SKIPPED = 'skipped'
124        FAILED = 'failed'
125
126    def __init__(self, status, *messages):
127        self._status = status
128        self._messages = list(messages)
129
130    def ok(self):
131        return self._status in {_Result.Status.DONE, _Result.Status.SKIPPED}
132
133    def status_str(self):
134        return self._status
135
136    def messages(self):
137        return self._messages
138
139
140def _process_globs(globs):
141    unique_globs = []
142    for pat in globs:
143        if pat and pat not in unique_globs:
144            unique_globs.append(pat)
145
146    files = []
147    warnings = []
148    for pat in unique_globs:
149        if pat:
150            matches = glob.glob(pat)
151            if not matches:
152                warnings.append(
153                    'warning: pattern "{}" matched 0 files'.format(pat))
154            files.extend(matches)
155
156    if globs and not files:
157        warnings.append('warning: matched 0 total files')
158
159    return files, warnings
160
161
162def result_func(glob_warnings):
163    def result(status, *args):
164        return _Result(status, *([str(x) for x in glob_warnings] + list(args)))
165
166    return result
167
168
169class ConfigFileError(Exception):
170    pass
171
172
173# TODO(mohrr) remove disable=useless-object-inheritance once in Python 3.
174# pylint: disable=useless-object-inheritance
175# pylint: disable=too-many-instance-attributes
176# pylint: disable=too-many-arguments
177class EnvSetup(object):
178    """Run environment setup for Pigweed."""
179    def __init__(self, pw_root, cipd_cache_dir, shell_file, quiet, install_dir,
180                 use_pigweed_defaults, cipd_package_file, virtualenv_root,
181                 virtualenv_requirements, virtualenv_gn_target,
182                 virtualenv_gn_out_dir, cargo_package_file, enable_cargo,
183                 json_file, project_root, config_file):
184        self._env = environment.Environment()
185        self._project_root = project_root
186        self._pw_root = pw_root
187        self._setup_root = os.path.join(pw_root, 'pw_env_setup', 'py',
188                                        'pw_env_setup')
189        self._cipd_cache_dir = cipd_cache_dir
190        self._shell_file = shell_file
191        self._is_windows = os.name == 'nt'
192        self._quiet = quiet
193        self._install_dir = install_dir
194        self._virtualenv_root = (virtualenv_root
195                                 or os.path.join(install_dir, 'pigweed-venv'))
196
197        if os.path.isfile(shell_file):
198            os.unlink(shell_file)
199
200        if isinstance(self._pw_root, bytes) and bytes != str:
201            self._pw_root = self._pw_root.decode()
202
203        self._cipd_package_file = []
204        self._virtualenv_requirements = []
205        self._virtualenv_gn_targets = []
206        self._cargo_package_file = []
207        self._enable_cargo = enable_cargo
208
209        if config_file:
210            self._parse_config_file(config_file)
211
212        self._json_file = json_file
213
214        setup_root = os.path.join(pw_root, 'pw_env_setup', 'py',
215                                  'pw_env_setup')
216
217        # TODO(pwbug/67, pwbug/68) Investigate pulling these files into an
218        # oxidized env setup executable instead of referring to them in the
219        # source tree. Note that this could be error-prone because users expect
220        # changes to the files in the source tree to affect bootstrap.
221        if use_pigweed_defaults:
222            # If updating this section make sure to update
223            # $PW_ROOT/pw_env_setup/docs.rst as well.
224            self._cipd_package_file.append(
225                os.path.join(setup_root, 'cipd_setup', 'pigweed.json'))
226            self._cipd_package_file.append(
227                os.path.join(setup_root, 'cipd_setup', 'luci.json'))
228            # Only set if no other GN target is provided.
229            if not virtualenv_gn_target:
230                self._virtualenv_gn_targets.append(
231                    virtualenv_setup.GnTarget(
232                        '{}#pw_env_setup:python.install'.format(pw_root)))
233            self._cargo_package_file.append(
234                os.path.join(setup_root, 'cargo_setup', 'packages.txt'))
235
236        self._cipd_package_file.extend(cipd_package_file)
237        self._virtualenv_requirements.extend(virtualenv_requirements)
238        self._virtualenv_gn_targets.extend(virtualenv_gn_target)
239        self._virtualenv_gn_out_dir = virtualenv_gn_out_dir
240        self._cargo_package_file.extend(cargo_package_file)
241
242        self._env.set('PW_PROJECT_ROOT', project_root)
243        self._env.set('PW_ROOT', pw_root)
244        self._env.set('_PW_ACTUAL_ENVIRONMENT_ROOT', install_dir)
245        self._env.add_replacement('_PW_ACTUAL_ENVIRONMENT_ROOT', install_dir)
246        self._env.add_replacement('PW_ROOT', pw_root)
247
248    def _parse_config_file(self, config_file):
249        config = json.load(config_file)
250
251        self._cipd_package_file.extend(
252            os.path.join(self._project_root, x)
253            for x in config.pop('cipd_package_files', ()))
254
255        virtualenv = config.pop('virtualenv', {})
256
257        if virtualenv.get('gn_root'):
258            root = os.path.join(self._project_root, virtualenv.pop('gn_root'))
259        else:
260            root = self._project_root
261
262        for target in virtualenv.pop('gn_targets', ()):
263            self._virtualenv_gn_targets.append(
264                virtualenv_setup.GnTarget('{}#{}'.format(root, target)))
265
266        if virtualenv:
267            raise ConfigFileError(
268                'unrecognized option in {}: "virtualenv.{}"'.format(
269                    config_file.name, next(iter(virtualenv))))
270
271        if config:
272            raise ConfigFileError('unrecognized option in {}: "{}"'.format(
273                config_file.name, next(iter(config))))
274
275    def _log(self, *args, **kwargs):
276        # Not using logging module because it's awkward to flush a log handler.
277        if self._quiet:
278            return
279        flush = kwargs.pop('flush', False)
280        print(*args, **kwargs)
281        if flush:
282            sys.stdout.flush()
283
284    def setup(self):
285        """Runs each of the env_setup steps."""
286
287        if os.name == 'nt':
288            windows_env_start.print_banner(bootstrap=True, no_shell_file=False)
289        else:
290            enable_colors()
291
292        steps = [
293            ('CIPD package manager', self.cipd),
294            ('Python environment', self.virtualenv),
295            ('Host tools', self.host_tools),
296        ]
297
298        # TODO(pwbug/63): Add a Windows version of cargo to CIPD.
299        if not self._is_windows and self._enable_cargo:
300            steps.append(("Rust cargo", self.cargo))
301
302        if self._is_windows:
303            steps.append(("Windows scripts", self.win_scripts))
304
305        self._log(
306            Color.bold('Downloading and installing packages into local '
307                       'source directory:\n'))
308
309        max_name_len = max(len(name) for name, _ in steps)
310
311        self._env.comment('''
312This file is automatically generated. DO NOT EDIT!
313For details, see $PW_ROOT/pw_env_setup/py/pw_env_setup/env_setup.py and
314$PW_ROOT/pw_env_setup/py/pw_env_setup/environment.py.
315'''.strip())
316
317        if not self._is_windows:
318            self._env.comment('''
319For help debugging errors in this script, uncomment the next line.
320set -x
321Then use `set +x` to go back to normal.
322'''.strip())
323
324        self._env.echo(
325            Color.bold(
326                'Activating environment (setting environment variables):'))
327        self._env.echo('')
328
329        for name, step in steps:
330            self._log('  Setting up {name:.<{width}}...'.format(
331                name=name, width=max_name_len),
332                      end='',
333                      flush=True)
334            self._env.echo(
335                '  Setting environment variables for {name:.<{width}}...'.
336                format(name=name, width=max_name_len),
337                newline=False,
338            )
339
340            spin = spinner.Spinner()
341            with spin():
342                result = step(spin)
343
344            self._log(result.status_str())
345
346            self._env.echo(result.status_str())
347            for message in result.messages():
348                sys.stderr.write('{}\n'.format(message))
349                self._env.echo(message)
350
351            if not result.ok():
352                return -1
353
354        self._log('')
355        self._env.echo('')
356
357        self._env.finalize()
358
359        self._env.echo(Color.bold('Checking the environment:'))
360        self._env.echo()
361
362        self._env.doctor()
363        self._env.echo()
364
365        self._env.echo(
366            Color.bold('Environment looks good, you are ready to go!'))
367        self._env.echo()
368
369        with open(self._shell_file, 'w') as outs:
370            self._env.write(outs)
371
372        deactivate = os.path.join(
373            self._install_dir,
374            'deactivate{}'.format(os.path.splitext(self._shell_file)[1]))
375        with open(deactivate, 'w') as outs:
376            self._env.write_deactivate(outs)
377
378        config = {
379            # Skipping sysname and nodename in os.uname(). nodename could change
380            # based on the current network. sysname won't change, but is
381            # redundant because it's contained in release or version, and
382            # skipping it here simplifies logic.
383            'uname': ' '.join(getattr(os, 'uname', lambda: ())()[2:]),
384            'os': os.name,
385        }
386
387        with open(os.path.join(self._install_dir, 'config.json'), 'w') as outs:
388            outs.write(
389                json.dumps(config, indent=4, separators=(',', ': ')) + '\n')
390
391        if self._json_file is not None:
392            with open(self._json_file, 'w') as outs:
393                self._env.json(outs)
394
395        return 0
396
397    def cipd(self, spin):
398        install_dir = os.path.join(self._install_dir, 'cipd')
399
400        try:
401            cipd_client = cipd_wrapper.init(install_dir, silent=True)
402        except cipd_wrapper.UnsupportedPlatform as exc:
403            return result_func(('    {!r}'.format(exc), ))(
404                _Result.Status.SKIPPED,
405                '    abandoning CIPD setup',
406            )
407
408        package_files, glob_warnings = _process_globs(self._cipd_package_file)
409        result = result_func(glob_warnings)
410
411        if not package_files:
412            return result(_Result.Status.SKIPPED)
413
414        if not cipd_update.update(cipd=cipd_client,
415                                  root_install_dir=install_dir,
416                                  package_files=package_files,
417                                  cache_dir=self._cipd_cache_dir,
418                                  env_vars=self._env,
419                                  spin=spin):
420            return result(_Result.Status.FAILED)
421
422        return result(_Result.Status.DONE)
423
424    def virtualenv(self, unused_spin):
425        """Setup virtualenv."""
426
427        requirements, req_glob_warnings = _process_globs(
428            self._virtualenv_requirements)
429        result = result_func(req_glob_warnings)
430
431        orig_python3 = _which('python3')
432        with self._env():
433            new_python3 = _which('python3')
434
435        # There is an issue with the virtualenv module on Windows where it
436        # expects sys.executable to be called "python.exe" or it fails to
437        # properly execute. If we installed Python 3 in the CIPD step we need
438        # to address this. Detect if we did so and if so create a copy of
439        # python3.exe called python.exe so that virtualenv works.
440        if orig_python3 != new_python3 and self._is_windows:
441            python3_copy = os.path.join(os.path.dirname(new_python3),
442                                        'python.exe')
443            if not os.path.exists(python3_copy):
444                shutil.copyfile(new_python3, python3_copy)
445            new_python3 = python3_copy
446
447        if not requirements and not self._virtualenv_gn_targets:
448            return result(_Result.Status.SKIPPED)
449
450        if not virtualenv_setup.install(
451                project_root=self._project_root,
452                venv_path=self._virtualenv_root,
453                requirements=requirements,
454                gn_targets=self._virtualenv_gn_targets,
455                gn_out_dir=self._virtualenv_gn_out_dir,
456                python=new_python3,
457                env=self._env,
458        ):
459            return result(_Result.Status.FAILED)
460
461        return result(_Result.Status.DONE)
462
463    def host_tools(self, unused_spin):
464        # The host tools are grabbed from CIPD, at least initially. If the
465        # user has a current host build, that build will be used instead.
466        # TODO(mohrr) find a way to do stuff like this for all projects.
467        host_dir = os.path.join(self._pw_root, 'out', 'host')
468        self._env.prepend('PATH', os.path.join(host_dir, 'host_tools'))
469        return _Result(_Result.Status.DONE)
470
471    def win_scripts(self, unused_spin):
472        # These scripts act as a compatibility layer for windows.
473        env_setup_dir = os.path.join(self._pw_root, 'pw_env_setup')
474        self._env.prepend('PATH', os.path.join(env_setup_dir,
475                                               'windows_scripts'))
476        return _Result(_Result.Status.DONE)
477
478    def cargo(self, unused_spin):
479        install_dir = os.path.join(self._install_dir, 'cargo')
480
481        package_files, glob_warnings = _process_globs(self._cargo_package_file)
482        result = result_func(glob_warnings)
483
484        if not package_files:
485            return result(_Result.Status.SKIPPED)
486
487        if not cargo_setup.install(install_dir=install_dir,
488                                   package_files=package_files,
489                                   env=self._env):
490            return result(_Result.Status.FAILED)
491
492        return result(_Result.Status.DONE)
493
494
495def parse(argv=None):
496    """Parse command-line arguments."""
497    parser = argparse.ArgumentParser()
498
499    pw_root = os.environ.get('PW_ROOT', None)
500    if not pw_root:
501        try:
502            with open(os.devnull, 'w') as outs:
503                pw_root = subprocess.check_output(
504                    ['git', 'rev-parse', '--show-toplevel'],
505                    stderr=outs).strip()
506        except subprocess.CalledProcessError:
507            pw_root = None
508
509    parser.add_argument(
510        '--pw-root',
511        default=pw_root,
512        required=not pw_root,
513    )
514
515    project_root = os.environ.get('PW_PROJECT_ROOT', None) or pw_root
516
517    parser.add_argument(
518        '--project-root',
519        default=project_root,
520        required=not project_root,
521    )
522
523    parser.add_argument(
524        '--cipd-cache-dir',
525        default=os.environ.get('CIPD_CACHE_DIR',
526                               os.path.expanduser('~/.cipd-cache-dir')),
527    )
528
529    parser.add_argument(
530        '--shell-file',
531        help='Where to write the file for shells to source.',
532        required=True,
533    )
534
535    parser.add_argument(
536        '--quiet',
537        help='Reduce output.',
538        action='store_true',
539        default='PW_ENVSETUP_QUIET' in os.environ,
540    )
541
542    parser.add_argument(
543        '--install-dir',
544        help='Location to install environment.',
545        required=True,
546    )
547
548    parser.add_argument(
549        '--config-file',
550        help='JSON file describing CIPD and virtualenv requirements.',
551        type=argparse.FileType('r'),
552    )
553
554    parser.add_argument(
555        '--use-pigweed-defaults',
556        help='Use Pigweed default values in addition to the given environment '
557        'variables.',
558        action='store_true',
559    )
560
561    parser.add_argument(
562        '--cipd-package-file',
563        help='CIPD package file. JSON file consisting of a list of dicts with '
564        '"path" and "tags" keys, where "tags" a list of str.',
565        default=[],
566        action='append',
567    )
568
569    parser.add_argument(
570        '--virtualenv-requirements',
571        help='Pip requirements file. Compiled with pip-compile.',
572        default=[],
573        action='append',
574    )
575
576    parser.add_argument(
577        '--virtualenv-gn-target',
578        help=('GN targets that build and install Python packages. Format: '
579              'path/to/gn_root#target'),
580        default=[],
581        action='append',
582        type=virtualenv_setup.GnTarget,
583    )
584
585    parser.add_argument(
586        '--virtualenv-gn-out-dir',
587        help=('Output directory to use when building and installing Python '
588              'packages with GN; defaults to a unique path in the environment '
589              'directory.'))
590
591    parser.add_argument(
592        '--virtualenv-root',
593        help=('Root of virtualenv directory. Default: '
594              '<install_dir>/pigweed-venv'),
595        default=None,
596    )
597
598    parser.add_argument(
599        '--cargo-package-file',
600        help='Rust cargo packages to install. Lines with package name and '
601        'version separated by a space.',
602        default=[],
603        action='append',
604    )
605
606    parser.add_argument(
607        '--enable-cargo',
608        help='Enable cargo installation.',
609        action='store_true',
610    )
611
612    parser.add_argument(
613        '--json-file',
614        help='Dump environment variable operations to a JSON file.',
615        default=None,
616    )
617
618    args = parser.parse_args(argv)
619
620    others = (
621        'use_pigweed_defaults',
622        'cipd_package_file',
623        'virtualenv_requirements',
624        'virtualenv_gn_target',
625        'cargo_package_file',
626    )
627
628    one_required = others + ('config_file', )
629
630    if not any(getattr(args, x) for x in one_required):
631        parser.error('At least one of ({}) is required'.format(', '.join(
632            '"--{}"'.format(x.replace('_', '-')) for x in one_required)))
633
634    if args.config_file and any(getattr(args, x) for x in others):
635        parser.error('Cannot combine --config-file with any of {}'.format(
636            ', '.join('"--{}"'.format(x.replace('_', '-'))
637                      for x in one_required)))
638
639    return args
640
641
642def main():
643    try:
644        return EnvSetup(**vars(parse())).setup()
645    except subprocess.CalledProcessError as err:
646        print()
647        print(err.output)
648        raise
649
650
651if __name__ == '__main__':
652    sys.exit(main())
653