Skip to content

Commit 603bda9

Browse files
scripts: clean up restart-strategy and service help output
argparse rendered enum choices via str(), so --restart-strategy and --service showed '{RestartStrategy.ALL_AT_ONCE,...}' / '{Service.Core,...}' rather than the tokens users type. Add __str__ to both enums returning the accepted token (value / name), give the args a metavar plus a help line listing the options, and add a concise help formatter so options render as '-j, --service SERVICE' instead of repeating the metavar. Also fix restart_strategy_converter to catch ValueError (Enum value lookup raises ValueError, not KeyError) so invalid input gets the informative error instead of argparse's generic fallback. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 85fcb3f commit 603bda9

4 files changed

Lines changed: 76 additions & 6 deletions

File tree

scripts/prod/common_lib.py

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -156,6 +156,11 @@ class RestartStrategy(Enum):
156156
ONE_BY_ONE = "one_by_one"
157157
NO_RESTART = "no_restart"
158158

159+
def __str__(self) -> str:
160+
# The accepted CLI token is the value (e.g. "all_at_once"); use it so argparse choices and
161+
# error messages show what the user actually types rather than "RestartStrategy.ALL_AT_ONCE".
162+
return self.value
163+
159164

160165
def restart_strategy_converter(strategy_name: str) -> RestartStrategy:
161166
"""Convert string to RestartStrategy enum with informative error message"""
@@ -166,8 +171,9 @@ def restart_strategy_converter(strategy_name: str) -> RestartStrategy:
166171
strategy_name = strategy_name.lower()
167172

168173
try:
174+
# Looking an Enum up by value raises ValueError (not KeyError) when no member matches.
169175
return RestartStrategy(strategy_name)
170-
except KeyError:
176+
except ValueError:
171177
valid_strategies = ", ".join([strategy.value for strategy in RestartStrategy])
172178
raise argparse.ArgumentTypeError(
173179
f"Invalid restart strategy '{strategy_name}'. Valid options are: {valid_strategies}"
@@ -190,6 +196,11 @@ def __init__(self, config_map_name: str, pod_name: str) -> None:
190196
self.config_map_name = config_map_name
191197
self.pod_name = pod_name
192198

199+
def __str__(self) -> str:
200+
# The accepted CLI token is the member name (e.g. "Core"); use it so argparse choices and
201+
# error messages show what the user actually types rather than "Service.Core".
202+
return self.name
203+
193204

194205
class NamespaceAndInstructionArgs:
195206
def __init__(
Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
#!/usr/bin/env python3
2+
"""Tests that enum CLI tokens (used in argparse choices/help/errors) match accepted input."""
3+
4+
import argparse
5+
6+
import pytest
7+
from common_lib import RestartStrategy, Service, restart_strategy_converter
8+
9+
10+
def test_restart_strategy_str_is_the_accepted_token():
11+
assert str(RestartStrategy.ALL_AT_ONCE) == "all_at_once"
12+
assert str(RestartStrategy.ONE_BY_ONE) == "one_by_one"
13+
assert str(RestartStrategy.NO_RESTART) == "no_restart"
14+
# The string form round-trips back through the converter.
15+
for strategy in RestartStrategy:
16+
assert restart_strategy_converter(str(strategy)) is strategy
17+
18+
19+
def test_service_str_is_the_accepted_token():
20+
assert str(Service.Core) == "Core"
21+
assert str(Service.SierraCompiler) == "SierraCompiler"
22+
# The string form round-trips back through name lookup (how --service is parsed).
23+
for service in Service:
24+
assert Service[str(service)] is service
25+
26+
27+
def test_invalid_restart_strategy_raises_informative_error():
28+
# Enum value lookup raises ValueError; the converter must translate it to ArgumentTypeError
29+
# with the valid options, rather than letting argparse fall back to a generic message.
30+
with pytest.raises(argparse.ArgumentTypeError) as error_info:
31+
restart_strategy_converter("bogus")
32+
message = str(error_info.value)
33+
assert "all_at_once" in message
34+
assert "one_by_one" in message
35+
assert "no_restart" in message
36+
37+
38+
if __name__ == "__main__":
39+
raise SystemExit(pytest.main([__file__, "-v"]))

scripts/prod/update_config_and_restart_nodes.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -137,9 +137,11 @@ def main():
137137
"-j",
138138
"--service",
139139
type=service_type_converter,
140-
choices=list(Service),
141140
default=Service.Core,
142-
help="Service type to operate on; determines configmap and pod names (default: Core)",
141+
metavar="SERVICE",
142+
help="Service type to operate on; determines configmap and pod names. One of: "
143+
+ ", ".join(str(service) for service in Service)
144+
+ " (default: Core)",
143145
)
144146

145147
args_builder.add_argument(

scripts/prod/update_config_and_restart_nodes_lib.py

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,23 @@
2727
from restarter_lib import ServiceRestarter
2828

2929

30+
class _ConciseHelpFormatter(argparse.RawDescriptionHelpFormatter):
31+
"""Help formatter that shows the metavar once per option.
32+
33+
Renders "-j, --service SERVICE" instead of argparse's default
34+
"-j SERVICE, --service SERVICE". This matches argparse's own behavior in Python 3.13+; this
35+
formatter provides it on the older versions we run on. Inherits RawDescriptionHelpFormatter so
36+
the multi-line usage epilog is still printed verbatim.
37+
"""
38+
39+
def _format_action_invocation(self, action: argparse.Action) -> str:
40+
if not action.option_strings or action.nargs == 0:
41+
return super()._format_action_invocation(action)
42+
default_metavar = self._get_default_metavar_for_optional(action)
43+
args_string = self._format_args(action, default_metavar)
44+
return ", ".join(action.option_strings) + " " + args_string
45+
46+
3047
class ApolloArgsParserBuilder:
3148
"""Builder class for creating argument parsers with required flags and custom arguments."""
3249

@@ -40,7 +57,7 @@ def __init__(self, description: str, usage_example: str, include_restart_strateg
4057
self.usage_example = usage_example
4158
self.parser = argparse.ArgumentParser(
4259
description=description,
43-
formatter_class=argparse.RawDescriptionHelpFormatter,
60+
formatter_class=_ConciseHelpFormatter,
4461
epilog=usage_example,
4562
)
4663

@@ -105,9 +122,10 @@ def _add_common_flags(self, include_restart_strategy: bool):
105122
"-t",
106123
"--restart-strategy",
107124
type=restart_strategy_converter,
108-
choices=list(RestartStrategy),
109125
required=True,
110-
help="Strategy for restarting nodes",
126+
metavar="STRATEGY",
127+
help="Strategy for restarting nodes. One of: "
128+
+ ", ".join(str(strategy) for strategy in RestartStrategy),
111129
)
112130

113131
def add_argument(self, *args, **kwargs):

0 commit comments

Comments
 (0)