Support for Enum arguments in argparse

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

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

Solution 2:

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) to tuple(e.name for e in enum_type)
• value = self._enum(values) to value = self._enum[values]