Support for Enum arguments in argparse
PythonArgparsePython Problem Overview
Is there a better way of supporting Enums as types of argparse arguments than this pattern?
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()]
Python Solutions
Solution 1 - Python
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))
Solution 2 - Python
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)totuple(e.name for e in enum_type)value = self._enum(values)tovalue = self._enum[values]
Solution 3 - Python
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.
Solution 4 - Python
Here's the relevant bug/issue: http://bugs.python.org/issue25061
Add native enum support for argparse
I already wrote too much there. :)
Solution 5 - Python
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()