支持argpars中的枚举参数

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

我知道这是一个老问题，但是我遇到了同样的问题（Python2.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

这是对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)

---

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

Answer 3

下面是相关的bug/问题：http://bugs.python.org/issue25061

为argparse添加本机枚举支持

我在那里写得太多了。：）