1# Copyright (c) 2024 Basalte bv
2#
3# SPDX-License-Identifier: Apache-2.0
4
5import argparse
6import os
7import subprocess
8import sys
9import textwrap
10from itertools import chain
11from pathlib import Path
12
13from west.commands import WestCommand
14from zephyr_ext_common import ZEPHYR_BASE
15
16sys.path.append(os.fspath(Path(__file__).parent.parent))
17import zephyr_module
18
19
20def in_venv() -> bool:
21    return sys.prefix != sys.base_prefix
22
23
24class Packages(WestCommand):
25    def __init__(self):
26        super().__init__(
27            "packages",
28            "manage packages for Zephyr",
29            "List and Install packages for Zephyr and modules",
30            accepts_unknown_args=True,
31        )
32
33    def do_add_parser(self, parser_adder):
34        parser = parser_adder.add_parser(
35            self.name,
36            help=self.help,
37            description=self.description,
38            formatter_class=argparse.RawDescriptionHelpFormatter,
39            epilog=textwrap.dedent(
40                """
41            Listing packages:
42
43                Run 'west packages <manager>' to list all dependencies
44                available from a given package manager, already
45                installed and not. These can be filtered by module,
46                see 'west packages <manager> --help' for details.
47            """
48            ),
49        )
50
51        parser.add_argument(
52            "-m",
53            "--module",
54            action="append",
55            default=[],
56            dest="modules",
57            metavar="<module>",
58            help="Zephyr module to run the 'packages' command for. "
59            "Use 'zephyr' if the 'packages' command should run for Zephyr itself. "
60            "Option can be passed multiple times. "
61            "If this option is not given, the 'packages' command will run for Zephyr "
62            "and all modules.",
63        )
64
65        subparsers_gen = parser.add_subparsers(
66            metavar="<manager>",
67            dest="manager",
68            help="select a manager.",
69            required=True,
70        )
71
72        pip_parser = subparsers_gen.add_parser(
73            "pip",
74            help="manage pip packages",
75            formatter_class=argparse.RawDescriptionHelpFormatter,
76            epilog=textwrap.dedent(
77                """
78            Manage pip packages:
79
80                Run 'west packages pip' to print all requirement files needed by
81                Zephyr and modules.
82
83                The output is compatible with the requirements file format itself.
84            """
85            ),
86        )
87
88        pip_parser.add_argument(
89            "--install",
90            action="store_true",
91            help="Install pip requirements instead of listing them. "
92            "A single 'pip install' command is built and executed. "
93            "Additional pip arguments can be passed after a -- separator "
94            "from the original 'west packages pip --install' command. For example pass "
95            "'--dry-run' to pip not to actually install anything, but print what would be.",
96        )
97
98        pip_parser.add_argument(
99            "--ignore-venv-check",
100            action="store_true",
101            help="Ignore the virtual environment check. "
102            "This is useful when running 'west packages pip --install' "
103            "in a CI environment where the virtual environment is not set up.",
104        )
105
106        return parser
107
108    def do_run(self, args, unknown):
109        if len(unknown) > 0 and unknown[0] != "--":
110            self.die(
111                f'Unknown argument "{unknown[0]}"; '
112                'arguments for the manager should be passed after "--"'
113            )
114
115        # Store the zephyr modules for easier access
116        self.zephyr_modules = zephyr_module.parse_modules(ZEPHYR_BASE, self.manifest)
117
118        if args.modules:
119            # Check for unknown module names
120            module_names = [m.meta.get("name") for m in self.zephyr_modules]
121            module_names.append("zephyr")
122            for m in args.modules:
123                if m not in module_names:
124                    self.die(f'Unknown zephyr module "{m}"')
125
126        if args.manager == "pip":
127            return self.do_run_pip(args, unknown[1:])
128
129        # Unreachable but print an error message if an implementation is missing.
130        self.die(f'Unsupported package manager: "{args.manager}"')
131
132    def do_run_pip(self, args, manager_args):
133        requirements = []
134
135        if not args.modules or "zephyr" in args.modules:
136            requirements.append(ZEPHYR_BASE / "scripts/requirements.txt")
137
138        for module in self.zephyr_modules:
139            module_name = module.meta.get("name")
140            if args.modules and module_name not in args.modules:
141                if args.install:
142                    self.dbg(f"Skipping module {module_name}")
143                continue
144
145            # Get the optional pip section from the package managers
146            pip = module.meta.get("package-managers", {}).get("pip")
147            if pip is None:
148                if args.install:
149                    self.dbg(f"Nothing to install for {module_name}")
150                continue
151
152            # Add requirements files
153            requirements += [Path(module.project) / r for r in pip.get("requirement-files", [])]
154
155        if args.install:
156            if not in_venv() and not args.ignore_venv_check:
157                self.die("Running pip install outside of a virtual environment")
158
159            if len(requirements) > 0:
160                subprocess.check_call(
161                    [sys.executable, "-m", "pip", "install"]
162                    + list(chain.from_iterable([("-r", r) for r in requirements]))
163                    + manager_args
164                )
165            else:
166                self.inf("Nothing to install")
167            return
168
169        if len(manager_args) > 0:
170            self.die(f'west packages pip does not support unknown arguments: "{manager_args}"')
171
172        self.inf("\n".join([f"-r {r}" for r in requirements]))
173