• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2022 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"""Script that preprocesses a Python command then runs it.
15
16This script evaluates expressions in the Python command's arguments then invokes
17the command.
18"""
19
20import argparse
21import atexit
22import json
23import logging
24import os
25from pathlib import Path
26import platform
27import shlex
28import subprocess
29import sys
30import time
31from typing import List, Optional, Tuple
32
33try:
34    from pw_build import gn_resolver
35    from pw_build.python_package import load_packages
36except (ImportError, ModuleNotFoundError):
37    # Load from python_package from this directory if pw_build is not available.
38    from python_package import load_packages  # type: ignore
39    import gn_resolver  # type: ignore
40
41if sys.platform != 'win32':
42    import fcntl  # pylint: disable=import-error
43
44    # TODO(b/227670947): Support Windows.
45
46_LOG = logging.getLogger(__name__)
47_LOCK_ACQUISITION_TIMEOUT = 30 * 60  # 30 minutes in seconds
48
49# TODO(frolv): Remove these aliases once downstream projects are migrated.
50GnPaths = gn_resolver.GnPaths
51expand_expressions = gn_resolver.expand_expressions
52
53
54def _parse_args() -> argparse.Namespace:
55    """Parses arguments for this script, splitting out the command to run."""
56
57    parser = argparse.ArgumentParser(description=__doc__)
58    parser.add_argument(
59        '--gn-root',
60        type=Path,
61        required=True,
62        help=(
63            'Path to the root of the GN tree; '
64            'value of rebase_path("//", root_build_dir)'
65        ),
66    )
67    parser.add_argument(
68        '--current-path',
69        type=Path,
70        required=True,
71        help='Value of rebase_path(".", root_build_dir)',
72    )
73    parser.add_argument(
74        '--default-toolchain', required=True, help='Value of default_toolchain'
75    )
76    parser.add_argument(
77        '--current-toolchain', required=True, help='Value of current_toolchain'
78    )
79    parser.add_argument('--module', help='Run this module instead of a script')
80    parser.add_argument(
81        '--env',
82        action='append',
83        help='Environment variables to set as NAME=VALUE',
84    )
85    parser.add_argument(
86        '--touch',
87        type=Path,
88        help='File to touch after the command is run',
89    )
90    parser.add_argument(
91        '--capture-output',
92        action='store_true',
93        help='Capture subcommand output; display only on error',
94    )
95    parser.add_argument(
96        '--working-directory',
97        type=Path,
98        help='Change to this working directory before running the subcommand',
99    )
100    parser.add_argument(
101        '--python-dep-list-files',
102        nargs='+',
103        type=Path,
104        help='Paths to text files containing lists of Python package metadata '
105        'json files.',
106    )
107    parser.add_argument(
108        '--python-virtualenv-config',
109        type=Path,
110        help='Path to a virtualenv json config to use for this action.',
111    )
112    parser.add_argument(
113        '--command-launcher', help='Arguments to prepend to Python command'
114    )
115    parser.add_argument(
116        'original_cmd',
117        nargs=argparse.REMAINDER,
118        help='Python script with arguments to run',
119    )
120    parser.add_argument(
121        '--lockfile',
122        type=Path,
123        help=(
124            'Path to a pip lockfile. Any pip execution will acquire an '
125            'exclusive lock on it, any other module a shared lock.'
126        ),
127    )
128    return parser.parse_args()
129
130
131class LockAcquisitionTimeoutError(Exception):
132    """Raised on a timeout."""
133
134
135def acquire_lock(lockfile: Path, exclusive: bool):
136    """Attempts to acquire the lock.
137
138    Args:
139      lockfile: pathlib.Path to the lock.
140      exclusive: whether this needs to be an exclusive lock.
141
142    Raises:
143      LockAcquisitionTimeoutError: If the lock is not acquired after a
144        reasonable time.
145    """
146    if sys.platform == 'win32':
147        # No-op on Windows, which doesn't have POSIX file locking.
148        # TODO(b/227670947): Get this working on Windows, too.
149        return
150
151    start_time = time.monotonic()
152    if exclusive:
153        lock_type = fcntl.LOCK_EX  # type: ignore[name-defined]
154    else:
155        lock_type = fcntl.LOCK_SH  # type: ignore[name-defined]
156    fd = os.open(lockfile, os.O_RDWR | os.O_CREAT)
157
158    # Make sure we close the file when the process exits. If we manage to
159    # acquire the lock below, closing the file will release it.
160    def cleanup():
161        os.close(fd)
162
163    atexit.register(cleanup)
164
165    backoff = 1
166    while time.monotonic() - start_time < _LOCK_ACQUISITION_TIMEOUT:
167        try:
168            fcntl.flock(  # type: ignore[name-defined]
169                fd, lock_type | fcntl.LOCK_NB  # type: ignore[name-defined]
170            )
171            return  # Lock acquired!
172        except BlockingIOError:
173            pass  # Keep waiting.
174
175        time.sleep(backoff * 0.05)
176        backoff += 1
177
178    raise LockAcquisitionTimeoutError(
179        f"Failed to acquire lock {lockfile} in {_LOCK_ACQUISITION_TIMEOUT}"
180    )
181
182
183class MissingPythonDependency(Exception):
184    """An error occurred while processing a Python dependency."""
185
186
187def _load_virtualenv_config(json_file_path: Path) -> Tuple[str, str]:
188    with json_file_path.open() as json_fp:
189        json_dict = json.load(json_fp)
190    return json_dict.get('interpreter'), json_dict.get('path')
191
192
193def main(  # pylint: disable=too-many-arguments,too-many-branches,too-many-locals
194    gn_root: Path,
195    current_path: Path,
196    original_cmd: List[str],
197    default_toolchain: str,
198    current_toolchain: str,
199    module: Optional[str],
200    env: Optional[List[str]],
201    python_dep_list_files: List[Path],
202    python_virtualenv_config: Optional[Path],
203    capture_output: bool,
204    touch: Optional[Path],
205    working_directory: Optional[Path],
206    command_launcher: Optional[str],
207    lockfile: Optional[Path],
208) -> int:
209    """Script entry point."""
210
211    python_paths_list = []
212    if python_dep_list_files:
213        py_packages = load_packages(
214            python_dep_list_files,
215            # If this python_action has no gn python_deps this file will be
216            # empty.
217            ignore_missing=True,
218        )
219
220        for pkg in py_packages:
221            top_level_source_dir = pkg.package_dir
222            if not top_level_source_dir:
223                raise MissingPythonDependency(
224                    'Unable to find top level source dir for the Python '
225                    f'package "{pkg}"'
226                )
227            # Don't add this dir to the PYTHONPATH if no __init__.py exists.
228            init_py_files = top_level_source_dir.parent.glob('*/__init__.py')
229            if not any(init_py_files):
230                continue
231            python_paths_list.append(
232                gn_resolver.abspath(top_level_source_dir.parent)
233            )
234
235        # Sort the PYTHONPATH list, it will be in a different order each build.
236        python_paths_list = sorted(python_paths_list)
237
238    if not original_cmd or original_cmd[0] != '--':
239        _LOG.error('%s requires a command to run', sys.argv[0])
240        return 1
241
242    # GN build scripts are executed from the root build directory.
243    root_build_dir = gn_resolver.abspath(Path.cwd())
244
245    tool = current_toolchain if current_toolchain != default_toolchain else ''
246    paths = gn_resolver.GnPaths(
247        root=gn_resolver.abspath(gn_root),
248        build=root_build_dir,
249        cwd=gn_resolver.abspath(current_path),
250        toolchain=tool,
251    )
252
253    command = [sys.executable]
254
255    python_interpreter = None
256    python_virtualenv = None
257    if python_virtualenv_config:
258        python_interpreter, python_virtualenv = _load_virtualenv_config(
259            python_virtualenv_config
260        )
261
262    if python_interpreter is not None:
263        command = [str(root_build_dir / python_interpreter)]
264
265    if command_launcher is not None:
266        command = shlex.split(command_launcher) + command
267
268    if module is not None:
269        command += ['-m', module]
270
271    run_args: dict = dict()
272    # Always inherit the environtment by default. If PYTHONPATH or VIRTUALENV is
273    # set below then the environment vars must be copied in or subprocess.run
274    # will run with only the new updated variables.
275    run_args['env'] = os.environ.copy()
276
277    if env is not None:
278        environment = os.environ.copy()
279        environment.update((k, v) for k, v in (a.split('=', 1) for a in env))
280        run_args['env'] = environment
281
282    script_command = original_cmd[0]
283    if script_command == '--':
284        script_command = original_cmd[1]
285
286    is_pip_command = (
287        module == 'pip' or 'pip_install_python_deps.py' in script_command
288    )
289
290    existing_env = run_args['env'] if 'env' in run_args else os.environ.copy()
291    new_env = {}
292    if python_virtualenv:
293        new_env['VIRTUAL_ENV'] = str(root_build_dir / python_virtualenv)
294        bin_folder = 'Scripts' if platform.system() == 'Windows' else 'bin'
295        new_env['PATH'] = os.pathsep.join(
296            [
297                str(root_build_dir / python_virtualenv / bin_folder),
298                existing_env.get('PATH', ''),
299            ]
300        )
301
302    if python_virtualenv and python_paths_list and not is_pip_command:
303        python_path_prepend = os.pathsep.join(
304            str(p) for p in set(python_paths_list)
305        )
306
307        # Append the existing PYTHONPATH to the new one.
308        new_python_path = os.pathsep.join(
309            path_str
310            for path_str in [
311                python_path_prepend,
312                existing_env.get('PYTHONPATH', ''),
313            ]
314            if path_str
315        )
316
317        new_env['PYTHONPATH'] = new_python_path
318
319    if 'env' not in run_args:
320        run_args['env'] = {}
321    run_args['env'].update(new_env)
322
323    if capture_output:
324        # Combine stdout and stderr so that error messages are correctly
325        # interleaved with the rest of the output.
326        run_args['stdout'] = subprocess.PIPE
327        run_args['stderr'] = subprocess.STDOUT
328
329    # Build the command to run.
330    try:
331        for arg in original_cmd[1:]:
332            command += gn_resolver.expand_expressions(paths, arg)
333    except gn_resolver.ExpressionError as err:
334        _LOG.error('%s: %s', sys.argv[0], err)
335        return 1
336
337    if working_directory:
338        run_args['cwd'] = working_directory
339
340    # TODO(b/235239674): Deprecate the --lockfile option as part of the Python
341    # GN template refactor.
342    if lockfile:
343        try:
344            acquire_lock(lockfile, is_pip_command)
345        except LockAcquisitionTimeoutError as exception:
346            _LOG.error('%s', exception)
347            return 1
348
349    _LOG.debug('RUN %s', ' '.join(shlex.quote(arg) for arg in command))
350
351    completed_process = subprocess.run(command, **run_args)
352
353    if completed_process.returncode != 0:
354        _LOG.debug(
355            'Command failed; exit code: %d', completed_process.returncode
356        )
357        if capture_output:
358            sys.stdout.buffer.write(completed_process.stdout)
359    elif touch:
360        # If a stamp file is provided and the command executed successfully,
361        # touch the stamp file to indicate a successful run of the command.
362        touch = touch.resolve()
363        _LOG.debug('TOUCH %s', touch)
364
365        # Create the parent directory in case GN / Ninja hasn't created it.
366        touch.parent.mkdir(parents=True, exist_ok=True)
367        touch.touch()
368
369    return completed_process.returncode
370
371
372if __name__ == '__main__':
373    sys.exit(main(**vars(_parse_args())))
374