1# Copyright (c) 2024 Basalte bv
2#
3# SPDX-License-Identifier: Apache-2.0
4
5import argparse
6import os
7import platform
8import subprocess
9import sys
10import textwrap
11from itertools import chain
12from pathlib import Path, PureWindowsPath
13
14from west.commands import WestCommand
15from west.util import quote_sh_list
16from zephyr_ext_common import ZEPHYR_BASE
17
18sys.path.append(os.fspath(Path(__file__).parent.parent))
19import zephyr_module
20
21
22def in_venv() -> bool:
23    return sys.prefix != sys.base_prefix
24
25
26class Packages(WestCommand):
27    def __init__(self):
28        super().__init__(
29            "packages",
30            "manage packages for Zephyr",
31            "List and Install packages for Zephyr and modules",
32            accepts_unknown_args=True,
33        )
34
35    def do_add_parser(self, parser_adder):
36        parser = parser_adder.add_parser(
37            self.name,
38            help=self.help,
39            description=self.description,
40            formatter_class=argparse.RawDescriptionHelpFormatter,
41            epilog=textwrap.dedent(
42                """
43            Listing packages:
44
45                Run 'west packages <manager>' to list all dependencies
46                available from a given package manager, already
47                installed and not. These can be filtered by module,
48                see 'west packages <manager> --help' for details.
49            """
50            ),
51        )
52
53        parser.add_argument(
54            "-m",
55            "--module",
56            action="append",
57            default=[],
58            dest="modules",
59            metavar="<module>",
60            help="Zephyr module to run the 'packages' command for. "
61            "Use 'zephyr' if the 'packages' command should run for Zephyr itself. "
62            "Option can be passed multiple times. "
63            "If this option is not given, the 'packages' command will run for Zephyr "
64            "and all modules.",
65        )
66
67        subparsers_gen = parser.add_subparsers(
68            metavar="<manager>",
69            dest="manager",
70            help="select a manager.",
71            required=True,
72        )
73
74        pip_parser = subparsers_gen.add_parser(
75            "pip",
76            help="manage pip packages",
77            formatter_class=argparse.RawDescriptionHelpFormatter,
78            epilog=textwrap.dedent(
79                """
80            Manage pip packages:
81
82                Run 'west packages pip' to print all requirement files needed by
83                Zephyr and modules.
84
85                The output is compatible with the requirements file format itself.
86            """
87            ),
88        )
89
90        pip_parser.add_argument(
91            "--install",
92            action="store_true",
93            help="Install pip requirements instead of listing them. "
94            "A single 'pip install' command is built and executed. "
95            "Additional pip arguments can be passed after a -- separator "
96            "from the original 'west packages pip --install' command. For example pass "
97            "'--dry-run' to pip not to actually install anything, but print what would be.",
98        )
99
100        pip_parser.add_argument(
101            "--ignore-venv-check",
102            action="store_true",
103            help="Ignore the virtual environment check. "
104            "This is useful when running 'west packages pip --install' "
105            "in a CI environment where the virtual environment is not set up.",
106        )
107
108        return parser
109
110    def do_run(self, args, unknown):
111        if len(unknown) > 0 and unknown[0] != "--":
112            self.die(
113                f'Unknown argument "{unknown[0]}"; '
114                'arguments for the manager should be passed after "--"'
115            )
116
117        # Store the zephyr modules for easier access
118        self.zephyr_modules = zephyr_module.parse_modules(ZEPHYR_BASE, self.manifest)
119
120        if args.modules:
121            # Check for unknown module names
122            module_names = [m.meta.get("name") for m in self.zephyr_modules]
123            module_names.append("zephyr")
124            for m in args.modules:
125                if m not in module_names:
126                    self.die(f'Unknown zephyr module "{m}"')
127
128        if args.manager == "pip":
129            return self.do_run_pip(args, unknown[1:])
130
131        # Unreachable but print an error message if an implementation is missing.
132        self.die(f'Unsupported package manager: "{args.manager}"')
133
134    def do_run_pip(self, args, manager_args):
135        requirements = []
136
137        if not args.modules or "zephyr" in args.modules:
138            requirements.append(ZEPHYR_BASE / "scripts/requirements.txt")
139
140        for module in self.zephyr_modules:
141            module_name = module.meta.get("name")
142            if args.modules and module_name not in args.modules:
143                if args.install:
144                    self.dbg(f"Skipping module {module_name}")
145                continue
146
147            # Get the optional pip section from the package managers
148            pip = module.meta.get("package-managers", {}).get("pip")
149            if pip is None:
150                if args.install:
151                    self.dbg(f"Nothing to install for {module_name}")
152                continue
153
154            # Add requirements files
155            requirements += [Path(module.project) / r for r in pip.get("requirement-files", [])]
156
157        if args.install:
158            if not in_venv() and not args.ignore_venv_check:
159                self.die("Running pip install outside of a virtual environment")
160
161            if len(requirements) > 0:
162                cmd = [sys.executable, "-m", "pip", "install"]
163                cmd += chain.from_iterable([("-r", str(r)) for r in requirements])
164                cmd += manager_args
165                self.dbg(quote_sh_list(cmd))
166
167                # Use os.execv to execute a new program, replacing the current west process,
168                # this unloads all python modules first and allows for pip to update packages safely
169                if platform.system() != 'Windows':
170                    os.execv(cmd[0], cmd)
171
172                # Only reachable on Windows systems
173                # Windows does not really support os.execv:
174                # https://github.com/python/cpython/issues/63323
175                # https://github.com/python/cpython/issues/101191
176                # Warn the users about permission errors as those reported in:
177                # https://github.com/zephyrproject-rtos/zephyr/issues/100296
178                cmdscript = (
179                    PureWindowsPath(__file__).parents[1] / "utils" / "west-packages-pip-install.cmd"
180                )
181                self.wrn(
182                    "Updating packages on Windows with 'west packages pip --install', that are "
183                    "currently in use by west, results in permission errors. Leaving your "
184                    "environment with conflicting package versions. Recommended is to start with "
185                    "a new environment in that case.\n\n"
186                    "To avoid this using powershell run the following command instead:\n"
187                    f"{sys.executable} -m pip install @((west packages pip) -split ' ')\n\n"
188                    "Using cmd.exe execute the helper script:\n"
189                    f"cmd /c {cmdscript}\n\n"
190                    "Running 'west packages pip --install -- --dry-run' can provide information "
191                    "without actually updating the environment."
192                )
193                subprocess.check_call(cmd)
194            else:
195                self.inf("Nothing to install")
196            return
197
198        if len(manager_args) > 0:
199            self.die(f'west packages pip does not support unknown arguments: "{manager_args}"')
200
201        self.inf("\n".join([f"-r {r}" for r in requirements]))
202