Skip to content

Commit 0a2ea69

Browse files
committed
feat(espefuse): Add support for chaining commands with click parser
1 parent aa80001 commit 0a2ea69

File tree

3 files changed

+105
-4
lines changed

3 files changed

+105
-4
lines changed

espefuse/__init__.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,7 @@ def get_efuses(
6767

6868
@click.group(
6969
cls=Group,
70-
# chain=True, # allow using multiple commands in a single run
70+
chain=True, # allow using multiple commands in a single run
7171
no_args_is_help=True,
7272
context_settings=dict(help_option_names=["-h", "--help"], max_content_width=120),
7373
help=f"espefuse.py v{esptool.__version__} - ESP32xx eFuse get/set tool",

espefuse/cli_util.py

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,10 @@
44

55
from collections import namedtuple
66
from io import StringIO
7+
from typing import Any
78

89
import rich_click as click
10+
from click.parser import OptionParser, ParsingState, _unpack_args
911
from espefuse.efuse.base_operations import BaseCommands
1012
import esptool
1113
from esptool.cli_util import Group as EsptoolGroup
@@ -85,11 +87,110 @@ def get_command_class(chip_name: str) -> BaseCommands:
8587
}
8688

8789

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+
88187
class Group(EsptoolGroup):
89188
DEPRECATED_OPTIONS = {
90189
"--file_name": "--file-name",
91190
}
92191

192+
command_class = ChainingCommand
193+
93194
@staticmethod
94195
def _split_to_groups(args: list[str]) -> tuple[list[list[str]], list[str]]:
95196
"""

espefuse/efuse/base_operations.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -98,7 +98,7 @@ def type_cast_value(self, ctx: click.Context, value: list[str]) -> tuple[Any, ..
9898
ctx.exit()
9999

100100
# Check if we have more values than allowed by max_arity
101-
if len(value) > self.max_arity * self.type.arity:
101+
if self.max_arity is not None and len(value) > self.max_arity * self.type.arity:
102102
raise click.BadParameter(
103103
f"Expected at most {self.max_arity} groups ({self.type.arity} values "
104104
f"each), got {len(value)} (values: {value})"
@@ -294,8 +294,8 @@ def burn_bit_cli(block, bit_number, **kwargs):
294294
"split - each eFuse block is placed into its own file.",
295295
)
296296
@click.option(
297-
"--file_name",
298-
type=click.Path(exists=True, writable=True),
297+
"--file-name",
298+
type=click.File("w"),
299299
default=sys.stdout,
300300
help="The path to the file in which to save the dump, if not specified, "
301301
"output to the console.",

0 commit comments

Comments
 (0)