支持argparse中的Enum参数

Question

有没有比这种模式更好的方式支持Enums作为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干净地支持枚举的方法。

使用自定义操作的解决方案：

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 答案的改进。通过覆盖__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

这是一个简单的方法：

class Color(str, Enum):
    red = 'red'
    blue = 'blue'

parser = ArgumentParser()
parser.add_argument('color', type=Color)
args = parser.parse_args()
print('Your color was:', args.color)