Python 支持 argparse 中的枚举参数

Support for Enum arguments in argparse

提问by Andrzej Pronobis

Is there a better way of supporting Enums as types of argparse arguments than this pattern?

有没有比这种模式更好的支持枚举作为 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()]

回答by ron rothman

I see this is an old question, but I just came across the same problem (Python 2.7) and here's how I solved it:

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

Note that defining __str__is required to get ArgumentParser's help output to include the human readable (values) of Color.

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

Some sample invocations:

一些示例调用：

=> 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}

---

---

Since the OP's question specified integers as values, here is a slightly modified version that works in that case (using the enum names, rather than the values, as the command line args):

由于 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))

The only drawback there is that a bad parameter causes an ugly KeyError. That's easily solved by adding just a bit more code, converting the lambda into a proper function.

唯一的缺点是错误的参数会导致丑陋的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))

回答by hpaulj

Here's the relevant bug/issue: http://bugs.python.org/issue25061

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

Add native enum support for argparse

为 argparse 添加原生枚举支持

I already wrote too much there. :)

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

回答by David Lechner

This in an improvement on ron rothman's answer. By also overriding __repr__and changing to_stringa bit, we can get a better error message from argparsewhen the user enters a bad value.

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

In ron rothman's example, if we pass the color yellowas a command line argument, we get the following error:

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

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

With the improved code above, if we pass threeas a command line argument, we get:

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

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

---

---

IMHO, in the simple case of just converting the name of the enum members to lower case, the OP's method seems simpler. However, for more complex conversion cases, this could be useful.

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

回答by Tim

Just came across this issue also, however, all of the proposed solutions require changing the original Enum with extra methods.

刚刚也遇到了这个问题，但是，所有提出的解决方案都需要使用额外的方法更改原始 Enum。

Argparse does have a way to do this cleanly with actions.

Argparse 确实有一种方法可以通过动作干净利落地做到这一点。

My solution to this uses a reusable custom Action:

我对此的解决方案使用可重用的自定义操作：

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)

The advantages of this solution is that it is simple to use and requires no extra boiler plate on Enum definitions.

这种解决方案的优点是使用简单，不需要额外的 Enum 定义样板。

If you prefer to specify the enum by name change:

如果您更喜欢通过名称更改指定枚举：

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

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