Support for Enum arguments in argparse

Each Answer to this Q is separated by one/two green lines.

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

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

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]

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.

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

Add native enum support for argparse

I already wrote too much there. 🙂

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


The answers/resolutions are collected from stackoverflow, are licensed under cc by-sa 2.5 , cc by-sa 3.0 and cc by-sa 4.0 .