如何让argparse与Enum枚举及默认值实现良好配合？

问题背景 #

我们有一个自定义枚举类，希望将其作为argparse的命令行参数使用，要求：

• 命令行输入接受枚举名称的小写形式（如"one"、"two"）
• 无论是否传入参数，命名空间中都存储对应的枚举成员
• 帮助信息中显示友好的枚举名称（而非MyEnum.TWO这类形式）

原实现存在两个问题：

1. 默认值设为字符串时，命名空间中会保留字符串而非枚举成员
2. 默认值设为枚举成员时，帮助信息会显示MyEnum.TWO而非友好的"two"

解决方案 #

修改自定义EnumAction类，同时处理默认值的转换与帮助信息的友好显示：

from argparse import Action, ArgumentParser, Namespace
from collections.abc import Sequence
from enum import auto, Enum
from typing import Any


class EnumAction[T](Action):
    _enum: type[T]
    _enum_map: dict[str, T]
    _default_display: str | None = None

    def __init__(
        self,
        option_strings: Sequence[str],
        dest: str,
        nargs: int | str | None = None,
        const: Any = None,
        default: str | T | None = None,
        type: type[T] | None = None,
        choices: Sequence[Any] | None = None,
        required: bool = False,
        help: str | None = None,
        metavar: str | tuple[str, ...] | None = None,
    ) -> None:
        if type is None or not issubclass(type, Enum):
            raise TypeError("type must be an Enum class when using EnumAction")
        if choices is not None:
            raise ValueError("Can't specify choices when using EnumAction")

        self._enum = type
        self._enum_map = {e.name.lower(): e for e in self._enum}
        # 处理默认值：将字符串转成枚举成员，同时记录友好显示名称
        if isinstance(default, str):
            self._default_display = default.lower()
            default = self._enum_map[self._default_display]
        elif isinstance(default, Enum):
            self._default_display = default.name.lower()

        choices = tuple(self._enum_map.keys())
        # 替换help中的%(default)s为友好名称
        if help and self._default_display:
            help = help.replace("%(default)s", self._default_display)

        super().__init__(
            option_strings,
            dest,
            nargs,
            const,
            default,
            type=None,  # 禁用内置type转换，用自己的逻辑
            choices=choices,
            required=required,
            help=help,
            metavar=metavar,
        )

    def __call__(
        self,
        parser: ArgumentParser,
        namespace: Namespace,
        values: str | Sequence[Any] | None,
        option_string: str | None = None,
    ) -> None:
        if values is None:
            return
        # 将输入的字符串转成枚举成员
        setattr(namespace, self.dest, self._enum_map[values.lower()])


class MyEnum(Enum):
    ONE = auto()
    TWO = auto()
    THREE = auto()


# 使用示例
parser = ArgumentParser()
parser.add_argument(
    "--enum",
    action=EnumAction,
    type=MyEnum,
    default="two",  # 可以传字符串，也可以传MyEnum.TWO
    help="An enum value (Default: %(default)s)",
)

# 测试
print(parser.parse_args([]))  # 输出：Namespace(enum=<MyEnum.TWO: 2>)
print(parser.parse_args(["--enum", "one"]))  # 输出：Namespace(enum=<MyEnum.ONE: 1>)
parser.print_help()
# 帮助信息显示：--enum {one,two,three}  An enum value (Default: two)

关键改动说明 #

1. 默认值处理：在__init__中判断默认值类型，如果是字符串则转换为对应的枚举成员，同时记录其小写名称用于帮助信息；如果是枚举成员，则直接提取其小写名称。
2. 帮助信息替换：自动将help文本中的%(default)s替换为友好的枚举名称，无需手动修改。
3. 输入转换：在__call__中确保输入的字符串（无论大小写）都能正确转换为枚举成员。
4. 类型校验强化：更严格地校验type参数必须是Enum类。

内容的提问来源于stack exchange，提问作者Shai Avr