import argparse
from enum import Enum
from typing import Any
[docs]
class EnumAction(argparse.Action):
"""Argparse action for handling Enums
Allows to show help with the enum values as choices, while returning an Enum value
once the argument is parsed.
The implementation is copied from this answer on slackoverflow: https://stackoverflow.com/a/60750535
Examples
--------
>>> class MyEnum(Enum):
>>> Foo = "foo"
>>> Bar = "bar"
>>> parser = argparse.ArgumentParser()
>>> parser.add_argument('--my_enum', type=MyEnum, action=EnumAction)
>>> parser.print_help() # prints --my_enum {foo,bar}
"""
[docs]
def __init__(self, **kwargs):
# Pop off the enum type from kwargs
enum_type = kwargs.pop("type", None)
# Ensure an Enum subclass is provided
if enum_type is None:
raise ValueError("Argument `type` must be assigned to an Enum class when using EnumAction. It is not set.")
if not issubclass(enum_type, Enum):
raise TypeError(
"Argument `type` must be an Enum class when the action is `EnumAction`, got type {}".format(
str(type(enum_type))
)
)
# Generate choices from the Enum values
kwargs.setdefault("choices", tuple(e.value for e in enum_type))
# init the action
super(EnumAction, self).__init__(**kwargs)
# keep the enum type to generate Enum instances from their values when called
self.enum = enum_type
def __call__(
self, parser: argparse.ArgumentParser, namespace: argparse.Namespace, values: Any, option_string: str = None
):
"""Convert value back into a `self._enum` instance
Parameters
----------
parser : argparse.ArgumentParser
The parser instance invoking this action
namespace : argparse.Namespace
The namespace object that will be returned by the parser.
values : str
The command line argument string
option_string : str, optional
The option string that was used to invoke this action, by default None
"""
# Convert the value back to an enum instance and set it in the namespace
value = self.enum(values)
setattr(namespace, self.dest, value)