Support for Enum arguments in argparse

Question

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

Answer 1

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