|
4 | 4 |
|
5 | 5 | from collections import namedtuple
|
6 | 6 | from io import StringIO
|
| 7 | +from typing import Any |
7 | 8 |
|
8 | 9 | import rich_click as click
|
| 10 | +from click.parser import OptionParser, ParsingState, _unpack_args |
9 | 11 | from espefuse.efuse.base_operations import BaseCommands
|
10 | 12 | import esptool
|
11 | 13 | from esptool.cli_util import Group as EsptoolGroup
|
@@ -85,11 +87,110 @@ def get_command_class(chip_name: str) -> BaseCommands:
|
85 | 87 | }
|
86 | 88 |
|
87 | 89 |
|
| 90 | +class ChainParser(OptionParser): |
| 91 | + """ |
| 92 | + This is a modified version of the OptionParser class from click.parser. |
| 93 | + It allows for the processing of arguments and options in interspersed order |
| 94 | + together with chaining commands. |
| 95 | + """ |
| 96 | + |
| 97 | + def _process_args_for_options(self, state: ParsingState) -> None: |
| 98 | + while state.rargs: |
| 99 | + arg = state.rargs.pop(0) |
| 100 | + arglen = len(arg) |
| 101 | + # Double dashes always handled explicitly regardless of what |
| 102 | + # prefixes are valid. |
| 103 | + if arg == "--": |
| 104 | + return |
| 105 | + # if the argument is a command, stop parsing options |
| 106 | + elif arg.replace("_", "-") in SUPPORTED_COMMANDS: |
| 107 | + state.largs.append(arg) |
| 108 | + return |
| 109 | + elif arg[:1] in self._opt_prefixes and arglen > 1: |
| 110 | + self._process_opts(arg, state) |
| 111 | + elif self.allow_interspersed_args: |
| 112 | + state.largs.append(arg) |
| 113 | + else: |
| 114 | + state.rargs.insert(0, arg) |
| 115 | + return |
| 116 | + |
| 117 | + def _process_args_for_args(self, state: ParsingState) -> None: |
| 118 | + pargs, args = _unpack_args( |
| 119 | + state.largs + state.rargs, [x.nargs for x in self._args] |
| 120 | + ) |
| 121 | + |
| 122 | + # This check is required because of the way we modify nargs in ChainingCommand |
| 123 | + if len(pargs) > 0: |
| 124 | + for idx, arg in enumerate(self._args): |
| 125 | + arg.process(pargs[idx], state) |
| 126 | + |
| 127 | + state.largs = args |
| 128 | + state.rargs = [] |
| 129 | + |
| 130 | + |
| 131 | +class ChainingCommand(click.RichCommand, click.Command): |
| 132 | + def __init__(self, *args, **kwargs): |
| 133 | + super().__init__(*args, **kwargs) |
| 134 | + |
| 135 | + def _is_option(self, arg: str) -> bool: |
| 136 | + return arg.startswith("--") or arg.startswith("-") |
| 137 | + |
| 138 | + def invoke(self, ctx: click.Context) -> Any: |
| 139 | + log.print(f'\n=== Run "{self.name}" command ===') |
| 140 | + return super().invoke(ctx) |
| 141 | + |
| 142 | + def parse_args(self, ctx: click.Context, args: list[str]): |
| 143 | + # This is a hack to set nargs of the last argument to the number of arguments |
| 144 | + # that will be processed separately |
| 145 | + param_changed = None |
| 146 | + for idx, arg in enumerate(args): |
| 147 | + # command found in args or option found after argument |
| 148 | + if arg.replace("_", "-") in SUPPORTED_COMMANDS or ( |
| 149 | + self._is_option(arg) and idx > 0 |
| 150 | + ): |
| 151 | + arguments_count = sum( |
| 152 | + isinstance(param, click.Argument) for param in self.params |
| 153 | + ) |
| 154 | + for param in self.params: |
| 155 | + if param.nargs != -1: |
| 156 | + continue |
| 157 | + # set nargs of parameter to actual count of arguments and deduct |
| 158 | + # arguments_count as each argument will be processed separately, |
| 159 | + # we only care about the last one with nargs=-1 |
| 160 | + # at the end we add 1 to account for the processedargument itself |
| 161 | + # e.g. if we have burn-bit BLOCK2 1 2 3, we want to set nargs to 3, |
| 162 | + # so we need to account for BLOCK2 being processed separately |
| 163 | + param.nargs = args.index(arg) - arguments_count + 1 |
| 164 | + param_changed = param |
| 165 | + if param.nargs == 0 and param.required: |
| 166 | + raise click.UsageError( |
| 167 | + f"Command `{self.name}` requires the `{param.name}` " |
| 168 | + "argument." |
| 169 | + ) |
| 170 | + break |
| 171 | + break |
| 172 | + ret = super().parse_args(ctx, args) |
| 173 | + # restore nargs of the last argument to -1, in case it is going to be used again |
| 174 | + if param_changed is not None: |
| 175 | + param.nargs = -1 |
| 176 | + return ret |
| 177 | + |
| 178 | + def make_parser(self, ctx: click.Context) -> OptionParser: |
| 179 | + """Creates the underlying option parser for this command.""" |
| 180 | + parser = ChainParser(ctx) |
| 181 | + parser.allow_interspersed_args = True |
| 182 | + for param in self.get_params(ctx): |
| 183 | + param.add_to_parser(parser, ctx) |
| 184 | + return parser |
| 185 | + |
| 186 | + |
88 | 187 | class Group(EsptoolGroup):
|
89 | 188 | DEPRECATED_OPTIONS = {
|
90 | 189 | "--file_name": "--file-name",
|
91 | 190 | }
|
92 | 191 |
|
| 192 | + command_class = ChainingCommand |
| 193 | + |
93 | 194 | @staticmethod
|
94 | 195 | def _split_to_groups(args: list[str]) -> tuple[list[list[str]], list[str]]:
|
95 | 196 | """
|
|
0 commit comments