如何使 argparse 与枚举和默认值完美配合？

我有一个枚举：

from enum import auto, Enum


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

，我想将它用作 argparse 的参数。更具体地说，我想创建一个接受枚举名称之一 ("one", "two", "three") 并且可能具有默认值的参数，并且相应的枚举成员存储在命名空间中。换句话说，我想做这样的事情：

from argparse import ArgumentParser
from enum import auto, Enum


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


enum_map = {e.name.lower(): e for e in MyEnum}

parser = ArgumentParser()
parser.add_argument(
    "--enum",
    default="two",
    choices=tuple(enum_map.keys()),
    help="An enum value (Default: %(default)s)",
)
args = parser.parse_args()
args.enum = enum_map[args.enum]
print(args.enum)

这是一些样板代码，我想使用自定义操作消除它。我从 this SO 答案中得到灵感，并得到了以下示例：

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


class EnumAction[T](Action):
    _enum: type[T]
    _enum_map: dict[str, T]

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

        self._enum = type
        type = None
        self._enum_map = {e.name.lower(): e for e in self._enum}
        choices = tuple(self._enum_map.keys())
        super().__init__(
            option_strings,
            dest,
            nargs,
            const,
            default,
            type,
            choices,
            required,
            help,
            metavar,
        )

    def __call__(
        self,
        parser: ArgumentParser,
        namespace: Namespace,
        values: str | Sequence[Any] | None,
        option_string: str | None = None,
    ) -> None:
        setattr(namespace, self.dest, self._enum_map[values])


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


parser = ArgumentParser()
parser.add_argument(
    "--enum",
    action=EnumAction,
    type=MyEnum,
    default="two",
    help="An enum value (Default: %(default)s)",
)
print(parser.parse_args())

它工作得很好，除了一个问题：如果我在没有 --enum 参数的情况下调用程序，我会得到默认值| ||，但它没有被我的操作处理，所以我得到字符串本身而不是枚举成员： "two" 最简单的解决方法是将默认值设置为

$ python demo.py --enum one
Namespace(enum=<MyEnum.ONE: 1>)

$ python demo.py
Namespace(enum='two')

而不是 default=MyEnum.TWO 或使我的操作转换如果是字符串，则将 "two" 参数放入相应的枚举成员中。但是，这会导致枚举成员显示在不太可读的帮助消息中（这就是我将选项设置为成员名称而不是实际成员的原因）： default 如何获得枚举在传递参数和使用默认值时都在命名空间中添加成员，并在帮助消息中看到成员的名称而不是成员本身？

$ python demo.py -h
usage: demo.py [-h] [--enum {one,two,three}]

options:
  -h, --help            show this help message and exit
  --enum {one,two,three}
                        An enum value (Default: MyEnum.TWO)

How can I get an enum member in the namespace both when passing an argument and when using the default value, and see the member's name in the help message instead of the member itself?

可以通过在 __call__ 方法中处理默认值并使用 argparse 的 metavar 参数来自定义帮助消息来实现此目的。

以下是更新后的代码：

from argparse import Action, ArgumentParser
from enum import Enum, auto


class EnumAction(Action):
    """
    Argparse action for handling Enums
    """

    def __init__(self, **kwargs):
        # Pop enum type from kwargs, since we need to handle it ourselves
        self._enum = kwargs.pop("type", None)
        if self._enum is None:
            raise ValueError("type must be assigned an Enum when using EnumAction")
        if not issubclass(self._enum, Enum):
            raise TypeError("type must be an Enum when using EnumAction")

        # Generate choices and enum map
        self._enum_map = {e.name.lower(): e for e in self._enum}
        kwargs["choices"] = tuple(self._enum_map.keys())

        # Use enum names for metavar display in help message
        kwargs.setdefault("metavar", '{' + ','.join(self._enum_map.keys()) + '}')

        super().__init__(**kwargs)

    def __call__(self, parser, namespace, values, option_string=None):
        if values is not None:
            # Convert the input value to the enum member
            value = self._enum_map.get(values.lower())
            if value is None:
                raise ValueError(
                    f"Invalid choice: '{values}' (choose from {', '.join(self._enum_map.keys())})"
                )
            setattr(namespace, self.dest, value)
        else:
            # Use the default if provided, otherwise set to None
            setattr(namespace, self.dest, getattr(self, "default", None))


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


parser = ArgumentParser()
parser.add_argument(
    "--enum",
    action=EnumAction,
    type=MyEnum,
    default=MyEnum.TWO,  # Set default to enum member
    help="An enum value (Default: %(default)s)",
)
args = parser.parse_args()

print(args)
print(args.enum)

代码解释：

1. 处理默认值：
   • 我们在 EnumAction 的 __call__ 方法中检查 values 是否为 None 。如果是，则表示没有提供命令行参数，我们使用 getattr(self, "default", None) 获取默认值。
2. 自定义帮助消息：
   • 我们在 EnumAction 的 __init__ 方法中使用 kwargs.setdefault("metavar", ...) 来设置 metavar 参数。这将覆盖默认的元变量显示，并在帮助消息中使用枚举成员的名称。
3. 使用枚举成员作为默认值：
   • 在 add_argument 调用中，我们将 default 设置为 MyEnum.TWO ，这将确保默认情况下使用正确的枚举成员。

现在，当运行脚本时，无论是否提供 --enum 参数，都将获得预期的行为和帮助消息：

$ python demo.py --enum one
Namespace(enum=<MyEnum.ONE: 1>)
<MyEnum.ONE: 1>

$ python demo.py 
Namespace(enum=<MyEnum.TWO: 2>)
<MyEnum.TWO: 2>

$ python demo.py -h
usage: demo.py [-h] [--enum {one,two,three}]

options:
  -h, --help            show this help message and exit
  --enum {one,two,three}
                        An enum value (Default: MyEnum.TWO)