xref: /aosp_15_r20/external/pigweed/pw_env_setup/py/pw_env_setup/env_setup.py (revision 61c4878ac05f98d0ceed94b57d316916de578985)
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"""
21
22import argparse
23import copy
24import glob
25import inspect
26import json
27import os
28import shutil
29import subprocess
30import sys
31import time
32
33# If we're running oxidized, filesystem-centric import hacks won't work. In that
34# case, jump straight to the imports and assume oxidation brought in the deps.
35if not getattr(sys, 'oxidized', False):
36    old_sys_path = copy.deepcopy(sys.path)
37    filename = None
38    if hasattr(sys.modules[__name__], '__file__'):
39        filename = __file__
40    else:
41        # Try introspection in environments where __file__ is not populated.
42        frame = inspect.currentframe()
43        if frame is not None:
44            filename = inspect.getfile(frame)
45    # If none of our strategies worked, we're in a strange runtime environment.
46    # The imports are almost certainly going to fail.
47    if filename is None:
48        raise RuntimeError(
49            'Unable to locate pw_env_setup module; cannot continue.\n'
50            '\n'
51            'Try updating to one of the standard Python implemetations:\n'
52            '  https://www.python.org/downloads/'
53        )
54    sys.path = [
55        os.path.abspath(os.path.join(filename, os.path.pardir, os.path.pardir))
56    ]
57    import pw_env_setup  # pylint: disable=unused-import
58
59    sys.path = old_sys_path
60
61# pylint: disable=wrong-import-position
62from pw_env_setup.cipd_setup import update as cipd_update
63from pw_env_setup.cipd_setup import wrapper as cipd_wrapper
64from pw_env_setup.colors import Color, enable_colors
65from pw_env_setup import environment
66from pw_env_setup import spinner
67from pw_env_setup import virtualenv_setup
68from pw_env_setup import windows_env_start
69
70
71def _which(
72    executable, pathsep=os.pathsep, use_pathext=None, case_sensitive=None
73):
74    if use_pathext is None:
75        use_pathext = os.name == 'nt'
76    if case_sensitive is None:
77        case_sensitive = os.name != 'nt' and sys.platform != 'darwin'
78
79    if not case_sensitive:
80        executable = executable.lower()
81
82    exts = None
83    if use_pathext:
84        exts = frozenset(os.environ['PATHEXT'].split(pathsep))
85        if not case_sensitive:
86            exts = frozenset(x.lower() for x in exts)
87        if not exts:
88            raise ValueError('empty PATHEXT')
89
90    paths = os.environ['PATH'].split(pathsep)
91    for path in paths:
92        try:
93            entries = frozenset(os.listdir(path))
94            if not case_sensitive:
95                entries = frozenset(x.lower() for x in entries)
96        except OSError:
97            continue
98
99        if exts:
100            for ext in exts:
101                if executable + ext in entries:
102                    return os.path.join(path, executable + ext)
103        else:
104            if executable in entries:
105                return os.path.join(path, executable)
106
107    return None
108
109
110class _Result:
111    class Status:
112        DONE = 'done'
113        SKIPPED = 'skipped'
114        FAILED = 'failed'
115
116    def __init__(self, status, *messages):
117        self._status = status
118        self._messages = list(messages)
119
120    def ok(self):
121        return self._status in {_Result.Status.DONE, _Result.Status.SKIPPED}
122
123    def status_str(self, duration=None):
124        if not duration:
125            return self._status
126
127        duration_parts = []
128        if duration > 60:
129            minutes = int(duration // 60)
130            duration %= 60
131            duration_parts.append('{}m'.format(minutes))
132        duration_parts.append('{:.1f}s'.format(duration))
133        return '{} ({})'.format(self._status, ''.join(duration_parts))
134
135    def messages(self):
136        return self._messages
137
138
139class ConfigError(Exception):
140    pass
141
142
143def result_func(glob_warnings=()):
144    def result(status, *args):
145        return _Result(status, *([str(x) for x in glob_warnings] + list(args)))
146
147    return result
148
149
150class ConfigFileError(Exception):
151    pass
152
153
154class MissingSubmodulesError(Exception):
155    pass
156
157
158def _assert_sequence(value):
159    assert isinstance(value, (list, tuple))
160    return value
161
162
163# pylint: disable=too-many-instance-attributes
164# pylint: disable=too-many-arguments
165class EnvSetup:
166    """Run environment setup for Pigweed."""
167
168    def __init__(
169        self,
170        pw_root,
171        cipd_cache_dir,
172        shell_file,
173        quiet,
174        install_dir,
175        strict,
176        virtualenv_gn_out_dir,
177        json_file,
178        project_root,
179        config_file,
180        use_existing_cipd,
181        check_submodules,
182        use_pinned_pip_packages,
183        cipd_only,
184        trust_cipd_hash,
185        additional_cipd_file,
186        disable_rosetta,
187    ):
188        self._env = environment.Environment()
189        self._project_root = project_root
190        self._pw_root = pw_root
191        self._setup_root = os.path.join(
192            pw_root, 'pw_env_setup', 'py', 'pw_env_setup'
193        )
194        self._cipd_cache_dir = cipd_cache_dir
195        self._shell_file = shell_file
196        self._env._shell_file = shell_file
197        self._is_windows = os.name == 'nt'
198        self._quiet = quiet
199        self._install_dir = install_dir
200        self._virtualenv_root = os.path.join(self._install_dir, 'pigweed-venv')
201        self._strict = strict
202        self._cipd_only = cipd_only
203        self._trust_cipd_hash = trust_cipd_hash
204        self._additional_cipd_file = additional_cipd_file
205        self._disable_rosetta = disable_rosetta
206
207        if os.path.isfile(shell_file):
208            os.unlink(shell_file)
209
210        if isinstance(self._pw_root, bytes) and bytes != str:
211            self._pw_root = self._pw_root.decode()
212
213        self._cipd_package_file = []
214        self._project_actions = []
215        self._virtualenv_requirements = []
216        self._virtualenv_constraints = []
217        self._virtualenv_gn_targets = []
218        self._virtualenv_gn_args = []
219        self._virtualenv_pip_install_disable_cache = False
220        self._virtualenv_pip_install_find_links = []
221        self._virtualenv_pip_install_offline = False
222        self._virtualenv_pip_install_require_hashes = False
223        self._use_pinned_pip_packages = use_pinned_pip_packages
224        self._optional_submodules = []
225        self._required_submodules = []
226        self._virtualenv_system_packages = False
227        self._pw_packages = []
228        self._root_variable = None
229
230        self._check_submodules = check_submodules
231
232        self._json_file = json_file
233        self._gni_file = None
234
235        self._config_file_name = config_file
236        self._env.set(
237            '_PW_ENVIRONMENT_CONFIG_FILE', os.path.abspath(config_file)
238        )
239        if config_file:
240            self._parse_config_file(config_file)
241
242        self._check_submodule_presence()
243
244        self._use_existing_cipd = use_existing_cipd
245        self._virtualenv_gn_out_dir = virtualenv_gn_out_dir
246
247        if self._root_variable:
248            self._env.set(self._root_variable, project_root, deactivate=False)
249        self._env.set('PW_PROJECT_ROOT', project_root, deactivate=False)
250        self._env.set('PW_ROOT', pw_root, deactivate=False)
251        self._env.set('_PW_ACTUAL_ENVIRONMENT_ROOT', install_dir)
252        self._env.set('VIRTUAL_ENV', self._virtualenv_root)
253        self._env.add_replacement('_PW_ACTUAL_ENVIRONMENT_ROOT', install_dir)
254        self._env.add_replacement('PW_ROOT', pw_root)
255
256    def _process_globs(self, globs):
257        unique_globs = []
258        for pat in globs:
259            if pat and pat not in unique_globs:
260                unique_globs.append(pat)
261
262        files = []
263        warnings = []
264        for pat in unique_globs:
265            if pat:
266                matches = glob.glob(pat)
267                if not matches:
268                    warning = 'pattern "{}" matched 0 files'.format(pat)
269                    warnings.append('warning: {}'.format(warning))
270                    if self._strict:
271                        raise ConfigError(warning)
272
273                files.extend(matches)
274
275        if globs and not files:
276            warnings.append('warning: matched 0 total files')
277            if self._strict:
278                raise ConfigError('matched 0 total files')
279
280        return files, warnings
281
282    def _parse_config_file(self, config_file):
283        # This should use pw_env_setup.config_file instead.
284        with open(config_file, 'r') as ins:
285            config = json.load(ins)
286
287            # While transitioning, allow environment config to be at the top of
288            # the JSON file or at '.["pw"]["pw_env_setup"]'.
289            config = config.get('pw', config)
290            config = config.get('pw_env_setup', config)
291
292        self._root_variable = config.pop('root_variable', None)
293
294        # This variable is not used by env setup since we already have it.
295        # However, other tools may use it, so we double-check that it's correct.
296        pigweed_root = os.path.join(
297            self._project_root,
298            config.pop('relative_pigweed_root', self._pw_root),
299        )
300        if os.path.abspath(self._pw_root) != os.path.abspath(pigweed_root):
301            raise ValueError(
302                'expected Pigweed root {!r} in config but found {!r}'.format(
303                    os.path.relpath(self._pw_root, self._project_root),
304                    os.path.relpath(pigweed_root, self._project_root),
305                )
306            )
307
308        rosetta = config.pop('rosetta', 'allow')
309        if rosetta not in ('never', 'allow', 'force'):
310            raise ValueError(rosetta)
311        self._rosetta = rosetta in ('allow', 'force')
312        if self._disable_rosetta:
313            self._rosetta = False
314        self._env.set('_PW_ROSETTA', str(int(self._rosetta)))
315
316        if 'json_file' in config:
317            self._json_file = config.pop('json_file')
318
319        self._gni_file = config.pop('gni_file', None)
320
321        self._optional_submodules.extend(
322            _assert_sequence(config.pop('optional_submodules', ()))
323        )
324        self._required_submodules.extend(
325            _assert_sequence(config.pop('required_submodules', ()))
326        )
327
328        if self._optional_submodules and self._required_submodules:
329            raise ValueError(
330                '{} contains both "optional_submodules" and '
331                '"required_submodules", but these options are mutually '
332                'exclusive'.format(self._config_file_name)
333            )
334
335        self._cipd_package_file.extend(
336            os.path.join(self._project_root, x)
337            for x in _assert_sequence(config.pop('cipd_package_files', ()))
338        )
339        self._cipd_package_file.extend(
340            os.path.join(self._project_root, x)
341            for x in self._additional_cipd_file or ()
342        )
343
344        for action in config.pop('project_actions', {}):
345            # We can add a 'phase' option in the future if we end up needing to
346            # support project actions at more than one point in the setup flow.
347            self._project_actions.append(
348                (action['import_path'], action['module_name'])
349            )
350
351        for pkg in _assert_sequence(config.pop('pw_packages', ())):
352            self._pw_packages.append(pkg)
353
354        virtualenv = config.pop('virtualenv', {})
355
356        if virtualenv.get('gn_root'):
357            root = os.path.join(self._project_root, virtualenv.pop('gn_root'))
358        else:
359            root = self._project_root
360
361        for target in _assert_sequence(virtualenv.pop('gn_targets', ())):
362            self._virtualenv_gn_targets.append(
363                virtualenv_setup.GnTarget('{}#{}'.format(root, target))
364            )
365
366        self._virtualenv_gn_args = _assert_sequence(
367            virtualenv.pop('gn_args', ())
368        )
369
370        self._virtualenv_system_packages = virtualenv.pop(
371            'system_packages', False
372        )
373
374        for req_txt in _assert_sequence(virtualenv.pop('requirements', ())):
375            self._virtualenv_requirements.append(
376                os.path.join(self._project_root, req_txt)
377            )
378
379        for constraint_txt in _assert_sequence(
380            virtualenv.pop('constraints', ())
381        ):
382            self._virtualenv_constraints.append(
383                os.path.join(self._project_root, constraint_txt)
384            )
385
386        for pip_cache_dir in _assert_sequence(
387            virtualenv.pop('pip_install_find_links', ())
388        ):
389            self._virtualenv_pip_install_find_links.append(pip_cache_dir)
390
391        self._virtualenv_pip_install_disable_cache = virtualenv.pop(
392            'pip_install_disable_cache', False
393        )
394        self._virtualenv_pip_install_offline = virtualenv.pop(
395            'pip_install_offline', False
396        )
397        self._virtualenv_pip_install_require_hashes = virtualenv.pop(
398            'pip_install_require_hashes', False
399        )
400
401        if virtualenv:
402            raise ConfigFileError(
403                'unrecognized option in {}: "virtualenv.{}"'.format(
404                    self._config_file_name, next(iter(virtualenv))
405                )
406            )
407
408        if config:
409            raise ConfigFileError(
410                'unrecognized option in {}: "{}"'.format(
411                    self._config_file_name, next(iter(config))
412                )
413            )
414
415    def _check_submodule_presence(self):
416        uninitialized = set()
417
418        # If there's no `.git` file or directory, then we are not in
419        # a git repo and must skip the git-submodule check.
420        if not os.path.exists(os.path.join(self._project_root, '.git')):
421            return
422
423        if not self._check_submodules:
424            return
425
426        cmd = ['git', 'submodule', 'status', '--recursive']
427
428        for line in subprocess.check_output(
429            cmd, cwd=self._project_root
430        ).splitlines():
431            if isinstance(line, bytes):
432                line = line.decode()
433            # Anything but an initial '-' means the submodule is initialized.
434            if not line.startswith('-'):
435                continue
436            uninitialized.add(line.split()[1])
437
438        missing = uninitialized - set(self._optional_submodules)
439        if self._required_submodules:
440            missing = set(self._required_submodules) & uninitialized
441
442        if missing:
443            print(
444                'Not all submodules are initialized. Please run the '
445                'following commands.',
446                file=sys.stderr,
447            )
448            print('', file=sys.stderr)
449
450            for miss in sorted(missing):
451                print(
452                    '    git submodule update --init {}'.format(miss),
453                    file=sys.stderr,
454                )
455            print('', file=sys.stderr)
456
457            if self._required_submodules:
458                print(
459                    'If these submodules are not required, remove them from '
460                    'the "required_submodules"',
461                    file=sys.stderr,
462                )
463
464            else:
465                print(
466                    'If these submodules are not required, add them to the '
467                    '"optional_submodules"',
468                    file=sys.stderr,
469                )
470
471            print('list in the environment config JSON file:', file=sys.stderr)
472            print('    {}'.format(self._config_file_name), file=sys.stderr)
473            print('', file=sys.stderr)
474
475            raise MissingSubmodulesError(', '.join(sorted(missing)))
476
477    def _write_gni_file(self):
478        if self._cipd_only:
479            return
480
481        gni_file = os.path.join(
482            self._project_root, 'build_overrides', 'pigweed_environment.gni'
483        )
484        if self._gni_file:
485            gni_file = os.path.join(self._project_root, self._gni_file)
486
487        with open(gni_file, 'w') as outs:
488            self._env.gni(outs, self._project_root, gni_file)
489        shutil.copy(gni_file, os.path.join(self._install_dir, 'logs'))
490
491    def _log(self, *args, **kwargs):
492        # Not using logging module because it's awkward to flush a log handler.
493        if self._quiet:
494            return
495        flush = kwargs.pop('flush', False)
496        print(*args, **kwargs)
497        if flush:
498            sys.stdout.flush()
499
500    def setup(self):
501        """Runs each of the env_setup steps."""
502
503        if os.name == 'nt':
504            windows_env_start.print_banner(bootstrap=True, no_shell_file=False)
505        else:
506            enable_colors()
507
508        steps = [
509            ('CIPD package manager', self.cipd),
510            ('Project actions', self.project_actions),
511            ('Python environment', self.virtualenv),
512            ('pw packages', self.pw_package),
513            ('Host tools', self.host_tools),
514        ]
515
516        if self._is_windows:
517            steps.append(("Windows scripts", self.win_scripts))
518
519        if self._cipd_only:
520            steps = [('CIPD package manager', self.cipd)]
521
522        self._log(
523            Color.bold(
524                'Downloading and installing packages into local '
525                'source directory:\n'
526            )
527        )
528
529        max_name_len = max(len(name) for name, _ in steps)
530
531        self._env.comment(
532            '''
533This file is automatically generated. DO NOT EDIT!
534For details, see $PW_ROOT/pw_env_setup/py/pw_env_setup/env_setup.py and
535$PW_ROOT/pw_env_setup/py/pw_env_setup/environment.py.
536'''.strip()
537        )
538
539        if not self._is_windows:
540            self._env.comment(
541                '''
542For help debugging errors in this script, uncomment the next line.
543set -x
544Then use `set +x` to go back to normal.
545'''.strip()
546            )
547
548        self._env.echo(
549            Color.bold(
550                'Activating environment (setting environment variables):'
551            )
552        )
553        self._env.echo('')
554
555        for name, step in steps:
556            self._log(
557                '  Setting up {name:.<{width}}...'.format(
558                    name=name, width=max_name_len
559                ),
560                end='',
561                flush=True,
562            )
563            self._env.echo(
564                '  Setting environment variables for '
565                '{name:.<{width}}...'.format(name=name, width=max_name_len),
566                newline=False,
567            )
568
569            start = time.time()
570            spin = spinner.Spinner(self._quiet)
571            with spin():
572                result = step(spin)
573            stop = time.time()
574
575            self._log(result.status_str(stop - start))
576
577            self._env.echo(result.status_str())
578            for message in result.messages():
579                sys.stderr.write('{}\n'.format(message))
580                self._env.echo(message)
581
582            if not result.ok():
583                return -1
584
585            # Log the environment state at the end of each step for debugging.
586            log_dir = os.path.join(self._install_dir, 'logs')
587            if not os.path.isdir(log_dir):
588                os.makedirs(log_dir)
589            actions_json = os.path.join(
590                log_dir, 'post-{}.json'.format(name.replace(' ', '_'))
591            )
592            with open(actions_json, 'w') as outs:
593                self._env.json(outs)
594
595            # This file needs to be written after the CIPD step and before the
596            # Python virtualenv step. It also needs to be rewritten after the
597            # Python virtualenv step, so it's easiest to just write it after
598            # every step.
599            self._write_gni_file()
600
601        # Only write stuff for GitHub Actions once, at the end.
602        if 'GITHUB_ACTIONS' in os.environ:
603            self._env.github(self._install_dir)
604
605        self._log('')
606        self._env.echo('')
607
608        self._env.finalize()
609
610        self._env.echo(Color.bold('Checking the environment:'))
611        self._env.echo()
612
613        self._env.doctor()
614        self._env.echo()
615
616        self._env.echo(
617            Color.bold('Environment looks good, you are ready to go!')
618        )
619        self._env.echo()
620
621        # Don't write new files if all we did was update CIPD packages.
622        if self._cipd_only:
623            return 0
624
625        with open(self._shell_file, 'w') as outs:
626            self._env.write(outs, shell_file=self._shell_file)
627
628        deactivate = os.path.join(
629            self._install_dir,
630            'deactivate{}'.format(os.path.splitext(self._shell_file)[1]),
631        )
632        with open(deactivate, 'w') as outs:
633            self._env.write_deactivate(outs, shell_file=deactivate)
634
635        config = {
636            # Skipping sysname and nodename in os.uname(). nodename could change
637            # based on the current network. sysname won't change, but is
638            # redundant because it's contained in release or version, and
639            # skipping it here simplifies logic.
640            'uname': ' '.join(getattr(os, 'uname', lambda: ())()[2:]),
641            'os': os.name,
642        }
643
644        with open(os.path.join(self._install_dir, 'config.json'), 'w') as outs:
645            outs.write(
646                json.dumps(config, indent=4, separators=(',', ': ')) + '\n'
647            )
648
649        json_file = self._json_file or os.path.join(
650            self._install_dir, 'actions.json'
651        )
652        with open(json_file, 'w') as outs:
653            self._env.json(outs)
654
655        return 0
656
657    def cipd(self, spin):
658        """Set up cipd and install cipd packages."""
659
660        install_dir = os.path.join(self._install_dir, 'cipd')
661
662        # There's no way to get to the UnsupportedPlatform exception if this
663        # flag is set, but this flag should only be set in LUCI builds which
664        # will always have CIPD.
665        if self._use_existing_cipd:
666            cipd_client = 'cipd'
667
668        else:
669            try:
670                cipd_client = cipd_wrapper.init(
671                    install_dir,
672                    silent=True,
673                    rosetta=self._rosetta,
674                )
675            except cipd_wrapper.UnsupportedPlatform as exc:
676                return result_func(('    {!r}'.format(exc),))(
677                    _Result.Status.SKIPPED,
678                    '    abandoning CIPD setup',
679                )
680
681        package_files, glob_warnings = self._process_globs(
682            self._cipd_package_file
683        )
684        result = result_func(glob_warnings)
685
686        if not package_files:
687            return result(_Result.Status.SKIPPED)
688
689        if not cipd_update.update(
690            cipd=cipd_client,
691            root_install_dir=install_dir,
692            package_files=package_files,
693            cache_dir=self._cipd_cache_dir,
694            env_vars=self._env,
695            rosetta=self._rosetta,
696            spin=spin,
697            trust_hash=self._trust_cipd_hash,
698        ):
699            return result(_Result.Status.FAILED)
700
701        return result(_Result.Status.DONE)
702
703    def project_actions(self, unused_spin):
704        """Perform project install actions.
705
706        This is effectively a limited plugin system for performing
707        project-specific actions (e.g. fetching tools) after CIPD but before
708        virtualenv setup.
709        """
710        result = result_func()
711
712        if not self._project_actions:
713            return result(_Result.Status.SKIPPED)
714
715        if sys.version_info[0] < 3:
716            raise ValueError(
717                'Project Actions require Python 3 or higher. '
718                'The current python version is %s' % sys.version_info
719            )
720
721        # Once Keir okays removing 2.7 support for env_setup, move this import
722        # to the main list of imports at the top of the file.
723        import importlib  # pylint: disable=import-outside-toplevel
724
725        for import_path, module_name in self._project_actions:
726            full_import_path = os.path.join(self._project_root, import_path)
727            sys.path.append(full_import_path)
728            mod = importlib.import_module(module_name)
729            mod.run_action(env=self._env)
730
731        return result(_Result.Status.DONE)
732
733    def virtualenv(self, unused_spin):
734        """Setup virtualenv."""
735
736        requirements, req_glob_warnings = self._process_globs(
737            self._virtualenv_requirements
738        )
739
740        constraints, constraint_glob_warnings = self._process_globs(
741            self._virtualenv_constraints
742        )
743
744        result = result_func(req_glob_warnings + constraint_glob_warnings)
745
746        orig_python3 = _which('python3')
747        with self._env():
748            new_python3 = _which('python3')
749
750        # There is an issue with the virtualenv module on Windows where it
751        # expects sys.executable to be called "python.exe" or it fails to
752        # properly execute. If we installed Python 3 in the CIPD step we need
753        # to address this. Detect if we did so and if so create a copy of
754        # python3.exe called python.exe so that virtualenv works.
755        if orig_python3 != new_python3 and self._is_windows:
756            python3_copy = os.path.join(
757                os.path.dirname(new_python3), 'python.exe'
758            )
759            if not os.path.exists(python3_copy):
760                shutil.copyfile(new_python3, python3_copy)
761            new_python3 = python3_copy
762
763        if not requirements and not self._virtualenv_gn_targets:
764            return result(_Result.Status.SKIPPED)
765
766        if not virtualenv_setup.install(
767            project_root=self._project_root,
768            venv_path=self._virtualenv_root,
769            requirements=requirements,
770            constraints=constraints,
771            pip_install_find_links=self._virtualenv_pip_install_find_links,
772            pip_install_offline=self._virtualenv_pip_install_offline,
773            pip_install_require_hashes=(
774                self._virtualenv_pip_install_require_hashes
775            ),
776            pip_install_disable_cache=(
777                self._virtualenv_pip_install_disable_cache
778            ),
779            gn_args=self._virtualenv_gn_args,
780            gn_targets=self._virtualenv_gn_targets,
781            gn_out_dir=self._virtualenv_gn_out_dir,
782            python=new_python3,
783            env=self._env,
784            system_packages=self._virtualenv_system_packages,
785            use_pinned_pip_packages=self._use_pinned_pip_packages,
786        ):
787            return result(_Result.Status.FAILED)
788
789        return result(_Result.Status.DONE)
790
791    def pw_package(self, unused_spin):
792        """Install "default" pw packages."""
793
794        result = result_func()
795
796        pkg_dir = os.path.join(self._install_dir, 'packages')
797        self._env.set('PW_PACKAGE_ROOT', pkg_dir)
798
799        if not os.path.isdir(pkg_dir):
800            os.makedirs(pkg_dir)
801
802        if not self._pw_packages:
803            return result(_Result.Status.SKIPPED)
804
805        for pkg in self._pw_packages:
806            print('installing {}'.format(pkg))
807            cmd = ['pw', 'package', 'install', pkg]
808
809            log = os.path.join(pkg_dir, '{}.log'.format(pkg))
810            try:
811                with open(log, 'w') as outs, self._env():
812                    print(*cmd, file=outs)
813                    subprocess.check_call(
814                        cmd,
815                        cwd=self._project_root,
816                        stdout=outs,
817                        stderr=subprocess.STDOUT,
818                    )
819            except subprocess.CalledProcessError:
820                with open(log, 'r') as ins:
821                    sys.stderr.write(ins.read())
822                    raise
823
824        return result(_Result.Status.DONE)
825
826    def host_tools(self, unused_spin):
827        # The host tools are grabbed from CIPD, at least initially. If the
828        # user has a current host build, that build will be used instead.
829        # TODO(mohrr) find a way to do stuff like this for all projects.
830        host_dir = os.path.join(self._pw_root, 'out', 'host')
831        self._env.prepend('PATH', os.path.join(host_dir, 'host_tools'))
832        return _Result(_Result.Status.DONE)
833
834    def win_scripts(self, unused_spin):
835        # These scripts act as a compatibility layer for windows.
836        env_setup_dir = os.path.join(self._pw_root, 'pw_env_setup')
837        self._env.prepend(
838            'PATH', os.path.join(env_setup_dir, 'windows_scripts')
839        )
840        return _Result(_Result.Status.DONE)
841
842
843def parse(argv=None):
844    """Parse command-line arguments."""
845    parser = argparse.ArgumentParser(prog="python -m pw_env_setup.env_setup")
846
847    pw_root = os.environ.get('PW_ROOT', None)
848    if not pw_root:
849        try:
850            with open(os.devnull, 'w') as outs:
851                pw_root = subprocess.check_output(
852                    ['git', 'rev-parse', '--show-toplevel'], stderr=outs
853                ).strip()
854        except subprocess.CalledProcessError:
855            pw_root = None
856
857    parser.add_argument(
858        '--pw-root',
859        default=pw_root,
860        required=not pw_root,
861    )
862
863    project_root = os.environ.get('PW_PROJECT_ROOT', None) or pw_root
864
865    parser.add_argument(
866        '--project-root',
867        default=project_root,
868        required=not project_root,
869    )
870
871    default_cipd_cache_dir = os.environ.get(
872        'CIPD_CACHE_DIR', os.path.expanduser('~/.cipd-cache-dir')
873    )
874    if 'PW_NO_CIPD_CACHE_DIR' in os.environ:
875        default_cipd_cache_dir = None
876
877    parser.add_argument('--cipd-cache-dir', default=default_cipd_cache_dir)
878
879    parser.add_argument(
880        '--no-cipd-cache-dir',
881        action='store_const',
882        const=None,
883        dest='cipd_cache_dir',
884    )
885
886    parser.add_argument(
887        '--trust-cipd-hash',
888        action='store_true',
889        help='Only run the cipd executable if the ensure file or command-line '
890        'has changed. Defaults to false since files could have been deleted '
891        'from the installation directory and cipd would add them back.',
892    )
893
894    parser.add_argument(
895        '--shell-file',
896        help='Where to write the file for shells to source.',
897        required=True,
898    )
899
900    parser.add_argument(
901        '--quiet',
902        help='Reduce output.',
903        action='store_true',
904        default='PW_ENVSETUP_QUIET' in os.environ,
905    )
906
907    parser.add_argument(
908        '--install-dir',
909        help='Location to install environment.',
910        required=True,
911    )
912
913    parser.add_argument(
914        '--config-file',
915        help='Path to pigweed.json file.',
916        default=os.path.join(project_root, 'pigweed.json'),
917    )
918
919    parser.add_argument(
920        '--additional-cipd-file',
921        help=(
922            'Path to additional CIPD files, in addition to those referenced by '
923            'the --config-file file.'
924        ),
925        action='append',
926    )
927
928    parser.add_argument(
929        '--virtualenv-gn-out-dir',
930        help=(
931            'Output directory to use when building and installing Python '
932            'packages with GN; defaults to a unique path in the environment '
933            'directory.'
934        ),
935    )
936
937    parser.add_argument('--json-file', help=argparse.SUPPRESS, default=None)
938
939    parser.add_argument(
940        '--use-existing-cipd',
941        help='Use cipd executable from the environment instead of fetching it.',
942        action='store_true',
943    )
944
945    parser.add_argument(
946        '--strict',
947        help='Fail if there are any warnings.',
948        action='store_true',
949    )
950
951    parser.add_argument(
952        '--unpin-pip-packages',
953        dest='use_pinned_pip_packages',
954        help='Do not use pins of pip packages.',
955        action='store_false',
956    )
957
958    parser.add_argument(
959        '--cipd-only',
960        help='Skip non-CIPD steps.',
961        action='store_true',
962    )
963
964    parser.add_argument(
965        '--skip-submodule-check',
966        help='Skip checking for submodule presence.',
967        dest='check_submodules',
968        action='store_false',
969    )
970
971    parser.add_argument(
972        '--disable-rosetta',
973        help=(
974            "Disable Rosetta on ARM Macs, regardless of what's in "
975            'pigweed.json.'
976        ),
977        action='store_true',
978    )
979
980    args = parser.parse_args(argv)
981
982    return args
983
984
985def main():
986    try:
987        return EnvSetup(**vars(parse())).setup()
988    except subprocess.CalledProcessError as err:
989        print()
990        print(err.output)
991        raise
992
993
994if __name__ == '__main__':
995    sys.exit(main())
996