python - 支持 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))