File size: 2,739 Bytes
d06fa57 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | # Copyright (c) 2024-present, Yumeow. Licensed under the MIT License.
import argparse
def add_minus_flags(parser: argparse.ArgumentParser):
"""
自动为 argparse.ArgumentParser 中含有下划线的参数添加别名。例如：
--fix_existing -> 添加别名 --fix-existing
--augment_OD_num -> 添加别名 --augment-OD-num
"""
for action in parser._actions:
# 仅处理可选参数（以 '-' 开头的参数）
if not action.option_strings:
continue
aliases_to_add = []
for opt in action.option_strings:
if opt.startswith('--') and '_' in opt:
new_alias = opt.replace('_', '-')
if new_alias not in action.option_strings:
aliases_to_add.append(new_alias)
for alias in aliases_to_add:
action.option_strings.append(alias)
parser._option_string_actions[alias] = action
return parser
def add_negation_flags(parser: argparse.ArgumentParser):
"""
自动为 parser 中的 store_true 参数添加对应的 --no-xxx 选项。
"""
args_to_add = []
# 建立现有 flag 索引，防止重复添加
existing_flags = set()
for action in parser._actions:
existing_flags.update(action.option_strings)
for action in parser._actions:
if isinstance(action, argparse._StoreTrueAction):
# 为当前 action 收集所有可能的否定别名
negation_aliases = []
for option in action.option_strings:
if not option.startswith('--'):
continue
# 生成否定形式：--use_new_model -> --no-use_new_model
raw_name = option.removeprefix('--')
neg_opt = f'--no-{raw_name}'
# 只有当这个 flag 还不存在时才添加
if neg_opt not in existing_flags:
negation_aliases.append(neg_opt)
# 如果生成了有效的否定别名，将它们打包准备添加
if negation_aliases:
# 选取第一个选项名作为帮助文档显示的名称
primary_name = action.option_strings[0].removeprefix('--')
args_to_add.append({
'options': negation_aliases, # 这里是一个列表
'dest': action.dest,
'action': 'store_false',
'default': action.default,
'help': f"Disable {primary_name}"
})
for arg in args_to_add:
parser.add_argument(*arg['options'], dest=arg['dest'], action=arg['action'], default=arg['default'], help=(
arg['help']
))
return parser
|