如何使 argparse 与枚举和默认值很好地配合工作？

Question

我有一个枚举：

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)

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

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)

如何在传递参数和使用默认值时在命名空间中获取枚举成员，并在帮助消息中查看成员的名称而不是成员本身？

Answer 1

我认为你把这个问题过于复杂化了。 您可以使用

type

参数将可调用对象传递给

add_argument

。 在这种情况下，您可以使用 lambda 函数从输入返回枚举对象。

from argparse import ArgumentParser
from enum import auto, Enum


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


parser = ArgumentParser()
parser.add_argument(
    '--enum', 
    default='two', 
    choices=[e.name.lower() for e in MyEnum], 
    type=lambda e:MyEnum[e.upper()],
    help="An enum value (Default: %(default)s)",
)