在 arparse 中支持 Enum 参数

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

Add native enum support for argparse

I already wrote too much there. :)

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:

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.

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):

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.

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

This in an improvement on ron rothman's answer. By also overriding __repr__ and changing to_string a bit, we can get a better error message from argparse when the user enters a bad value.

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 yellow as a command line argument, we get the following error:

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

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

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.

Just came across this issue also; however, all of the proposed solutions require adding new methods to the Enum definition.

argparse includes a way of supporting an enum cleanly using actions.

The solution using a custom Action:

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)

Usage

class Do(enum.Enum):
Foo = "foo"
Bar = "bar"




parser = argparse.ArgumentParser()
parser.add_argument('do', type=Do, action=EnumAction)

The advantages of this solution are that it will work with any Enum without requiring additional boilerplate code while remaining simple to use.

If you prefer to specify the enum by name change:

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

Building on the answer by @Tim here is an extension to use enumeration names instead of values and print pretty error messages:

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.name for e in enum_type))


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


self._enum = enum_type


def __call__(self,
parser: argparse.ArgumentParser,
namespace: argparse.Namespace,
value: Any,
option_string: str = None):
# Convert value back into an Enum
if isinstance(value, str):
value = self._enum[value]
setattr(namespace, self.dest, value)
elif value is None:
raise argparse.ArgumentTypeError(
f"You need to pass a value after {option_string}!")
else:
# A pretty invalid choice message will be generated by argparse
raise argparse.ArgumentTypeError()

Here's a simple way:

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)