支持 argparse 中的枚举参数

Question

有没有比这种模式更好的方法来支持枚举作为 argparse 参数的类型？

class SomeEnum(Enum):
    ONE = 1
    TWO = 2

parser.add_argument('some_val', type=str, default='one',
                    choices=[i.name.lower() for i in SomeEnum])
...
args.some_val = SomeEnum[args.some_val.upper()]

Answer 1

我看到这是一个老问题，但我刚刚遇到了同样的问题（Python 2.7），我是这样解决的：

from argparse import ArgumentParser
from enum import Enum

class Color(Enum):
    red = 'red'
    blue = 'blue'
    green = 'green'

    def __str__(self):
        return self.value

parser = ArgumentParser()
parser.add_argument('color', type=Color, choices=list(Color))

opts = parser.parse_args()
print 'your color was:', opts.color

请注意，需要定义 __str__ 以获取 ArgumentParser 的帮助输出以包含 Color 的人类可读（值）。

一些示例调用：

=> python enumtest.py blue
your color was: blue

=> python enumtest.py not-a-color
usage: enumtest.py [-h] {blue,green,red}
enumtest.py: error: argument color: invalid Color value: 'not-a-color'

=> python enumtest.py -h
usage: enumtest.py [-h] {blue,green,red}

positional arguments:
  {blue,green,red}

---

由于 OP 的问题将整数指定为值，因此这里有一个稍微修改过的版本，适用于这种情况（使用枚举名称，而不是值，作为命令行参数）：

class Color(Enum):
    red = 1
    blue = 2
    green = 3

    def __str__(self):
        return self.name

parser = ArgumentParser()
parser.add_argument('color', type=lambda color: Color[color], choices=list(Color))

唯一的缺点是错误的参数会导致丑陋的KeyError。这很容易解决，只需添加更多代码，将 lambda 转换为适当的函数。

class Color(Enum):
    red = 1
    blue = 2
    green = 3

    def __str__(self):
        return self.name

    @staticmethod
    def from_string(s):
        try:
            return Color[s]
        except KeyError:
            raise ValueError()

parser = ArgumentParser()
parser.add_argument('color', type=Color.from_string, choices=list(Color))

Answer 2

刚刚也遇到了这个问题；但是，所有提议的解决方案都需要在 Enum 定义中添加新方法。

argparse 包含一种使用 actions 干净地支持枚举的方法。

使用自定义Action的解决方案：

import argparse
import enum


class EnumAction(argparse.Action):
    """
    Argparse action for handling Enums
    """
    def __init__(self, **kwargs):
        # Pop off the type value
        enum_type = kwargs.pop("type", None)

        # Ensure an Enum subclass is provided
        if enum_type is None:
            raise ValueError("type must be assigned an Enum when using EnumAction")
        if not issubclass(enum_type, enum.Enum):
            raise TypeError("type must be an Enum when using EnumAction")

        # Generate choices from the Enum
        kwargs.setdefault("choices", tuple(e.value for e in enum_type))

        super(EnumAction, self).__init__(**kwargs)

        self._enum = enum_type

    def __call__(self, parser, namespace, values, option_string=None):
        # Convert value back into an Enum
        value = self._enum(values)
        setattr(namespace, self.dest, value)

用法

class Do(enum.Enum):
    Foo = "foo"
    Bar = "bar"


parser = argparse.ArgumentParser()
parser.add_argument('do', type=Do, action=EnumAction)

此解决方案的优点是它可以与任何 Enum 一起使用，无需额外的样板代码，同时保持简单易用。

如果您更喜欢通过name 指定枚举，请更改：

• tuple(e.value for e in enum) 到 tuple(e.name for e in enum_type)
• value = self._enum(values) 到 value = self._enum[values]

Answer 3

这是对ron rothman's answer 的改进。通过覆盖__repr__ 并稍微更改to_string，我们可以在用户输入错误值时从argparse 获得更好的错误消息。

import argparse
import enum


class SomeEnum(enum.IntEnum):
    ONE = 1
    TWO = 2

    # magic methods for argparse compatibility

    def __str__(self):
        return self.name.lower()

    def __repr__(self):
        return str(self)

    @staticmethod
    def argparse(s):
        try:
            return SomeEnum[s.upper()]
        except KeyError:
            return s


parser = argparse.ArgumentParser()
parser.add_argument('some_val', type=SomeEnum.argparse, choices=list(SomeEnum))
args = parser.parse_args()
print('success:', type(args.some_val), args.some_val)

在 ron rothman 的示例中，如果我们将颜色 yellow 作为命令行参数传递，则会收到以下错误：

demo.py: error: argument color: invalid from_string value: 'yellow'

使用上面改进的代码，如果我们将three 作为命令行参数传递，我们会得到：

demo.py: error: argument some_val: invalid choice: 'three' (choose from one, two)

---

恕我直言，在将枚举成员的名称转换为小写的简单情况下，OP 的方法似乎更简单。但是，对于更复杂的转换情况，这可能很有用。

Answer 4

这是相关的错误/问题：http://bugs.python.org/issue25061

为 argparse 添加原生枚举支持

我已经在那里写了太多了。 :)

Answer 5

基于@Tim 的答案，这里是使用枚举名称而不是值并打印漂亮错误消息的扩展：

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

    def __init__(self, **kwargs):
        # Pop off the type value
        enum_type = kwargs.pop("type", None)

        # Ensure an Enum subclass is provided
        if enum_type is None:
            raise ValueError(
                "type must be assigned an Enum when using EnumAction")
        if not issubclass(enum_type, enum.Enum):
            raise TypeError("type must be an Enum when using EnumAction")

        # Generate choices from the Enum
        kwargs.setdefault("choices", tuple(e.name for e in enum_type))

        super(EnumAction, self).__init__(**kwargs)

        self._enum = enum_type

    def __call__(self,
                 parser: argparse.ArgumentParser,
                 namespace: argparse.Namespace,
                 value: Any,
                 option_string: str = None):
        # Convert value back into an Enum
        if isinstance(value, str):
            value = self._enum[value]
            setattr(namespace, self.dest, value)
        elif value is None:
            raise argparse.ArgumentTypeError(
                f"You need to pass a value after {option_string}!")
        else:
            # A pretty invalid choice message will be generated by argparse
            raise argparse.ArgumentTypeError()