xref: /aosp_15_r20/external/pigweed/pw_cli/py/pw_cli/arguments.py (revision 61c4878ac05f98d0ceed94b57d316916de578985)
1# Copyright 2020 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"""Defines arguments for the pw command."""
15
16import argparse
17from dataclasses import dataclass
18from enum import Enum
19from functools import cached_property
20import logging
21from pathlib import Path
22import sys
23from typing import Any, NoReturn
24
25from pw_cli import argument_types, plugins
26from pw_cli.branding import banner
27import pw_cli.env
28
29_HELP_HEADER = '''The Pigweed command line interface (CLI).
30
31Example uses:
32    pw logdemo
33    pw --loglevel debug watch -C out
34'''
35
36
37def parse_args() -> argparse.Namespace:
38    return arg_parser().parse_args()
39
40
41class ShellCompletionFormat(Enum):
42    """Supported shell tab completion modes."""
43
44    BASH = 'bash'
45    FISH = 'fish'
46    ZSH = 'zsh'
47
48
49@dataclass(frozen=True)
50class ShellCompletion:
51    """Transforms argparse actions into bash, fish, zsh shell completions."""
52
53    action: argparse.Action
54    parser: argparse.ArgumentParser
55
56    @property
57    def option_strings(self) -> list[str]:
58        return list(self.action.option_strings)
59
60    @cached_property
61    def help(self) -> str:
62        return self.parser._get_formatter()._expand_help(  # pylint: disable=protected-access
63            self.action
64        )
65
66    @property
67    def choices(self) -> list[str]:
68        return list(self.action.choices) if self.action.choices else []
69
70    @property
71    def flag(self) -> bool:
72        return self.action.nargs == 0
73
74    @property
75    def default(self) -> Any:
76        return self.action.default
77
78    def bash_option(self, text: str) -> list[str]:
79        result: list[str] = []
80        for option_str in self.option_strings:
81            if option_str.startswith(text):
82                result.append(option_str)
83        return result
84
85    def zsh_option(self, text: str) -> list[str]:
86        result: list[str] = []
87        for option_str in self.option_strings:
88            if not option_str.startswith(text):
89                continue
90
91            short_and_long_opts = ' '.join(self.option_strings)
92            # '(-h --help)-h[Display help message and exit]'
93            # '(-h --help)--help[Display help message and exit]'
94            help_text = self.help if self.help else ''
95
96            state_str = ''
97            if not self.flag:
98                state_str = ': :->' + option_str
99
100            result.append(
101                f'({short_and_long_opts}){option_str}[{help_text}]'
102                f'{state_str}'
103            )
104        return result
105
106    def fish_option(self, text: str) -> list[str]:
107        result: list[str] = []
108        for option_str in self.option_strings:
109            if not option_str.startswith(text):
110                continue
111
112            output: list[str] = []
113            if option_str.startswith('--'):
114                output.append(f'--long-option\t{option_str.lstrip("-")}')
115            elif option_str.startswith('-'):
116                output.append(f'--short-option\t{option_str.lstrip("-")}')
117
118            if self.choices:
119                choice_str = " ".join(self.choices)
120                output.append('--exclusive')
121                output.append('--arguments')
122                output.append(f'(string split " " "{choice_str}")')
123            elif self.action.type == Path:
124                output.append('--require-parameter')
125                output.append('--force-files')
126
127            if self.help:
128                output.append(f'--description\t"{self.help}"')
129
130            result.append('\t'.join(output))
131        return result
132
133
134def get_options_and_help(
135    parser: argparse.ArgumentParser,
136) -> list[ShellCompletion]:
137    return list(
138        ShellCompletion(action=action, parser=parser)
139        for action in parser._actions  # pylint: disable=protected-access
140    )
141
142
143def print_completions_for_option(
144    parser: argparse.ArgumentParser,
145    text: str = '',
146    tab_completion_format: str = ShellCompletionFormat.BASH.value,
147) -> None:
148    matched_lines: list[str] = []
149    for completion in get_options_and_help(parser):
150        if tab_completion_format == ShellCompletionFormat.ZSH.value:
151            matched_lines.extend(completion.zsh_option(text))
152        if tab_completion_format == ShellCompletionFormat.FISH.value:
153            matched_lines.extend(completion.fish_option(text))
154        else:
155            matched_lines.extend(completion.bash_option(text))
156
157    for line in matched_lines:
158        print(line)
159
160
161def print_banner() -> None:
162    """Prints the PIGWEED (or project specific) banner to stderr."""
163    parsed_env = pw_cli.env.pigweed_environment()
164    if parsed_env.PW_ENVSETUP_NO_BANNER or parsed_env.PW_ENVSETUP_QUIET:
165        return
166    print(banner() + '\n', file=sys.stderr)
167
168
169def format_help(registry: plugins.Registry) -> str:
170    """Returns the pw help information as a string."""
171    return f'{arg_parser().format_help()}\n{registry.short_help()}'
172
173
174class _ArgumentParserWithBanner(argparse.ArgumentParser):
175    """Parser that the Pigweed banner when there are parsing errors."""
176
177    def error(self, message: str) -> NoReturn:
178        print_banner()
179        self.print_usage(sys.stderr)
180        self.exit(2, f'{self.prog}: error: {message}\n')
181
182
183def add_tab_complete_arguments(
184    parser: argparse.ArgumentParser,
185) -> argparse.ArgumentParser:
186    parser.add_argument(
187        '--tab-complete-option',
188        nargs='?',
189        help='Print tab completions for the supplied option text.',
190    )
191    parser.add_argument(
192        '--tab-complete-format',
193        choices=list(shell.value for shell in ShellCompletionFormat),
194        default='bash',
195        help='Output format for tab completion results.',
196    )
197    return parser
198
199
200def arg_parser() -> argparse.ArgumentParser:
201    """Creates an argument parser for the pw command."""
202    argparser = _ArgumentParserWithBanner(
203        prog='pw',
204        add_help=False,
205        description=_HELP_HEADER,
206        formatter_class=argparse.RawDescriptionHelpFormatter,
207    )
208
209    # Do not use the built-in help argument so that displaying the help info can
210    # be deferred until the pw plugins have been registered.
211    argparser.add_argument(
212        '-h',
213        '--help',
214        action='store_true',
215        help='Display this help message and exit',
216    )
217    argparser.add_argument(
218        '-C',
219        '--directory',
220        type=argument_types.directory,
221        default=Path.cwd(),
222        help='Change to this directory before doing anything',
223    )
224    argparser.add_argument(
225        '-l',
226        '--loglevel',
227        type=argument_types.log_level,
228        default=logging.INFO,
229        help='Set the log level (debug, info, warning, error, critical)',
230    )
231    argparser.add_argument(
232        '--debug-log',
233        type=Path,
234        help=(
235            'Additionally log to this file at debug level; does not affect '
236            'terminal output'
237        ),
238    )
239    argparser.add_argument(
240        '--banner',
241        action=argparse.BooleanOptionalAction,
242        default=True,
243        help='Whether to print the Pigweed banner',
244    )
245    argparser.add_argument(
246        '--tab-complete-command',
247        nargs='?',
248        help='Print tab completions for the supplied command text.',
249    )
250
251    default_analytics = None
252    if pw_cli.env.pigweed_environment().PW_DISABLE_CLI_ANALYTICS:
253        default_analytics = False
254    argparser.add_argument(
255        '--analytics',
256        action='store_true',
257        default=default_analytics,
258        help='Temporarily enable analytics collection.',
259    )
260    argparser.add_argument(
261        '--no-analytics',
262        action='store_false',
263        dest='analytics',
264        help='Temporarily disable analytics collection.',
265    )
266
267    argparser.add_argument(
268        'command',
269        nargs='?',
270        help='Which command to run; see supported commands below',
271    )
272    argparser.add_argument(
273        'plugin_args',
274        metavar='...',
275        nargs=argparse.REMAINDER,
276        help='Remaining arguments are forwarded to the command',
277    )
278    return add_tab_complete_arguments(argparser)
279