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的问题将整数指定为值，这里有一个稍微修改过的版本，在这种情况下有效（使用枚举名称，而不是值，作为命令行args）：

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提供了一种通过操作完全支持枚举的方法。

具有自定义操作的解决方案：

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

        # Ensure an Enum subclass is provided
        if enum is None:
            raise ValueError("type must be assigned an Enum when using EnumAction")
        if not issubclass(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))

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

        self._enum = enum

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

# Usage
class Do(Enum):
    Foo = "foo"
    Bar = "bar"

parser = 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)
enum = self._enum(values)至enum = self._enum[values]

Answer 3

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

为argparse添加本机枚举支持

我在那里写得太多了。 :)

Answer 4

这是对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的方法似乎更简单。但是，对于更复杂的转换情况，这可能会有用。

支持argparse中的Enum参数

4 个答案: