1#!/usr/bin/env python3
2#
3# Copyright (c) 2024 Raspberry Pi (Trading) Ltd.
4#
5# SPDX-License-Identifier: BSD-3-Clause
6#
7# Common helpers and variables shared across Bazel-related Python scripts.
8
9import argparse
10import logging
11import os
12from pathlib import Path
13import shlex
14import shutil
15import subprocess
16import sys
17
18
19_LOG = logging.getLogger(__file__)
20
21SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))
22
23SDK_ROOT = subprocess.run(
24    (
25        "git",
26        "rev-parse",
27        "--show-toplevel",
28    ),
29    cwd=SCRIPT_DIR,
30    text=True,
31    check=True,
32    capture_output=True,
33).stdout.strip()
34
35
36def parse_common_args():
37    parser = argparse.ArgumentParser()
38    add_common_args(parser)
39    return parser.parse_args()
40
41def add_common_args(parser):
42    parser.add_argument(
43        "--picotool-dir",
44        help="Use a local copy of Picotool rather than the dynamically fetching it",
45        default=None,
46        type=Path,
47    )
48
49def override_picotool_arg(picotool_dir):
50    return f"--override_module=picotool={picotool_dir.resolve()}"
51
52
53def bazel_command() -> str:
54    """Return the path to bazelisk or bazel."""
55    if shutil.which("bazelisk"):
56        return shutil.which("bazelisk")
57    if shutil.which("bazel"):
58        return "bazel"
59
60    raise FileNotFoundError(
61        "Cannot find 'bazel' or 'bazelisk' in the current system PATH"
62    )
63
64
65def run_bazel(args, check=False, **kwargs):
66    command = (
67        bazel_command(),
68        *args,
69    )
70    _LOG.info("Running Bazel command: %s", shlex.join(command))
71    proc = subprocess.run(
72        command,
73        cwd=SDK_ROOT,
74        **kwargs,
75    )
76    if proc.returncode != 0:
77        _LOG.error("Command invocation failed with return code %d!", proc.returncode)
78        _LOG.error(
79            "Failing command: %s",
80            " ".join(shlex.quote(str(arg)) for arg in args),
81        )
82        if kwargs.get("capture_output", False):
83            output = (
84                proc.stderr if isinstance(proc.stderr, str) else proc.stderr.decode()
85            )
86            _LOG.error(
87                "Output:\n%s",
88                output,
89            )
90        if check:
91            raise subprocess.CalledProcessError(
92                returncode=proc.returncode,
93                cmd=command,
94                output=proc.stdout,
95                stderr=proc.stderr,
96            )
97    return proc
98
99
100def print_to_stderr(*args, **kwargs):
101    print(*args, file=sys.stderr, **kwargs)
102
103
104def print_framed_string(s):
105    """Frames a string of text and prints it to highlight script steps."""
106    header_spacer = "#" * (len(s) + 12)
107    print_to_stderr(header_spacer)
108    print_to_stderr("###   " + s + "   ###")
109    print_to_stderr(header_spacer)
110
111
112def setup_logging():
113    log_levels = [
114        (logging.ERROR, "\x1b[31m[ERROR]\x1b[0m"),
115        (logging.WARNING, "\x1b[33m[WARNING]\x1b[0m"),
116        (logging.INFO, "\x1b[35m[INFO]\x1b[0m"),
117        (logging.DEBUG, "\x1b[34m[DEBUG]\x1b[0m"),
118    ]
119    for level, level_text in log_levels:
120        logging.addLevelName(level, level_text)
121    logging.basicConfig(format="%(levelname)s %(message)s", level=logging.DEBUG)
122