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