xref: /aosp_15_r20/external/pigweed/pw_build/py/pw_build/exec.py (revision 61c4878ac05f98d0ceed94b57d316916de578985) 
1# Copyright 2019 The Pigweed Authors
2#
3# Licensed under the Apache License, Version 2.0 (the "License"); you may not
4# use this file except in compliance with the License. You may obtain a copy of
5# the License at
6#
7#     https://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
11# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
12# License for the specific language governing permissions and limitations under
13# the License.
14"""Python wrapper that runs a program. For use in GN."""
15
16import argparse
17import logging
18import os
19import re
20import shlex
21import subprocess
22import sys
23import pathlib
24
25# Need to be able to run without pw_cli installed in the virtualenv.
26try:
27    import pw_cli.log
28except ImportError:
29    pass
30
31_LOG = logging.getLogger(__name__)
32
33
34def argument_parser(
35    parser: argparse.ArgumentParser | None = None,
36) -> argparse.ArgumentParser:
37    """Registers the script's arguments on an argument parser."""
38
39    if parser is None:
40        parser = argparse.ArgumentParser(description=__doc__)
41
42    parser.add_argument(
43        '--args-file',
44        type=argparse.FileType('r'),
45        help='File containing extra positional arguments to the program',
46    )
47    parser.add_argument(
48        '--capture-output',
49        action='store_true',
50        help='Hide output from the program unless it fails',
51    )
52    parser.add_argument(
53        '-e',
54        '--env',
55        action='append',
56        default=[],
57        help='key=value environment pair for the process',
58    )
59    parser.add_argument(
60        '--env-file',
61        type=argparse.FileType('r'),
62        help='File defining environment variables for the process',
63    )
64    parser.add_argument(
65        '--skip-empty-args',
66        action='store_true',
67        help='Don\'t run the program if --args-file is empty',
68    )
69    parser.add_argument(
70        '--target',
71        help='GN build target that runs the program',
72    )
73    parser.add_argument(
74        '--working-directory',
75        type=pathlib.Path,
76        help='Directory to execute program in',
77    )
78    parser.add_argument(
79        'command',
80        nargs=argparse.REMAINDER,
81        help='Program to run with arguments',
82    )
83
84    return parser
85
86
87_ENV_REGEX = re.compile(r'(\w+)(\+)?=(.+)')
88
89
90def apply_env_var(string: str, env: dict[str, str]) -> None:
91    """Update an environment map with provided a key-value pair.
92
93    Pairs are accepted in two forms:
94
95      KEY=value    sets environment variable "KEY" to "value"
96      KEY+=value   appends OS-specific PATH separator and "value" to
97                   environment variable "KEY"
98    """
99    result = _ENV_REGEX.search(string.strip())
100    if not result:
101        return
102
103    key, append, val = result.groups()
104    if append is not None:
105        curr = env.get(key)
106        val = f'{curr}{os.path.pathsep}{val}' if curr else val
107
108    env[key] = val
109
110
111def main() -> int:
112    """Runs a program specified by command-line arguments."""
113    args = argument_parser().parse_args()
114    if not args.command or args.command[0] != '--':
115        return 1
116
117    env = os.environ.copy()
118
119    # Command starts after the "--".
120    command = args.command[1:]
121    # command[0] is the invoker.prog from gn and gn will escape
122    # the various spaces in the command which means when argparse
123    # gets the string argparse believes this as a single argument
124    # and cannot correctly break the string into a list that
125    # subprocess can handle.  By splitting the first element
126    # in the command list, if there is a space, all of the
127    # command[0] elements will be made into a list and if not
128    # then split won't do everything and the old behavior
129    # will continue.
130    front_command = command[0].split(' ')
131    del command[0]
132    command = front_command + command
133    extra_kw_args = {}
134
135    if args.args_file is not None:
136        empty = True
137        for line in args.args_file:
138            empty = False
139            command.append(line.strip())
140
141        if args.skip_empty_args and empty:
142            return 0
143
144    if args.env_file is not None:
145        for line in args.env_file:
146            apply_env_var(line, env)
147
148    # Apply command-line overrides at a higher priority than the env file.
149    for string in args.env:
150        apply_env_var(string, env)
151
152    if args.capture_output:
153        extra_kw_args['stdout'] = subprocess.PIPE
154        extra_kw_args['stderr'] = subprocess.STDOUT
155
156    if args.working_directory:
157        extra_kw_args['cwd'] = args.working_directory
158
159    process = subprocess.run(command, env=env, **extra_kw_args)  # type: ignore
160
161    if process.returncode != 0 and args.capture_output:
162        _LOG.error('')
163        _LOG.error(
164            'Command failed with exit code %d in GN build.', process.returncode
165        )
166        _LOG.error('')
167        _LOG.error('Build target:')
168        _LOG.error('')
169        _LOG.error('  %s', args.target)
170        _LOG.error('')
171        _LOG.error('Full command:')
172        _LOG.error('')
173        _LOG.error('  %s', ' '.join(shlex.quote(arg) for arg in command))
174        _LOG.error('')
175        _LOG.error('Process output:')
176        print(flush=True)
177        sys.stdout.buffer.write(process.stdout)
178        print(flush=True)
179        _LOG.error('')
180
181    return process.returncode
182
183
184if __name__ == '__main__':
185    # If pw_cli is not yet installed in the virtualenv just skip it.
186    if 'pw_cli' in globals():
187        pw_cli.log.install()
188    sys.exit(main())
189