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        return parser
99
100    def do_run(self, args, unknown):
101        if len(unknown) > 0 and unknown[0] != "--":
102            self.die(
103                f'Unknown argument "{unknown[0]}"; '
104                'arguments for the manager should be passed after "--"'
105            )
106
107        # Store the zephyr modules for easier access
108        self.zephyr_modules = zephyr_module.parse_modules(ZEPHYR_BASE, self.manifest)
109
110        if args.modules:
111            # Check for unknown module names
112            module_names = [m.meta.get("name") for m in self.zephyr_modules]
113            module_names.append("zephyr")
114            for m in args.modules:
115                if m not in module_names:
116                    self.die(f'Unknown zephyr module "{m}"')
117
118        if args.manager == "pip":
119            return self.do_run_pip(args, unknown[1:])
120
121        # Unreachable but print an error message if an implementation is missing.
122        self.die(f'Unsupported package manager: "{args.manager}"')
123
124    def do_run_pip(self, args, manager_args):
125        requirements = []
126
127        if not args.modules or "zephyr" in args.modules:
128            requirements.append(ZEPHYR_BASE / "scripts/requirements.txt")
129
130        for module in self.zephyr_modules:
131            module_name = module.meta.get("name")
132            if args.modules and module_name not in args.modules:
133                if args.install:
134                    self.dbg(f"Skipping module {module_name}")
135                continue
136
137            # Get the optional pip section from the package managers
138            pip = module.meta.get("package-managers", {}).get("pip")
139            if pip is None:
140                if args.install:
141                    self.dbg(f"Nothing to install for {module_name}")
142                continue
143
144            # Add requirements files
145            requirements += [Path(module.project) / r for r in pip.get("requirement-files", [])]
146
147        if args.install:
148            if not in_venv():
149                self.die("Running pip install outside of a virtual environment")
150
151            if len(requirements) > 0:
152                subprocess.check_call(
153                    [sys.executable, "-m", "pip", "install"]
154                    + list(chain.from_iterable([("-r", r) for r in requirements]))
155                    + manager_args
156                )
157            else:
158                self.inf("Nothing to install")
159            return
160
161        if len(manager_args) > 0:
162            self.die(f'west packages pip does not support unknown arguments: "{manager_args}"')
163
164        self.inf("\n".join([f"-r {r}" for r in requirements]))
165