Поддержка аргументов Enum в argparse

Ответ 1

Я вижу, что это старый вопрос, но я столкнулся с одной и той же проблемой (Python 2.7), и вот как я ее решил:

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

Обратите внимание, что для определения ArgumentParser для вывода ArgumentParser требуется указать __str__, чтобы включить считываемый человеком (значения) Color.

Некоторые примеры запросов:

=> 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}

Поскольку вопрос ОП задает целые числа как значения, вот немного модифицированная версия, которая работает в этом случае (с использованием имен перечислений, а не значений, как указано в командной строке):

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

Единственный недостаток заключается в том, что плохой параметр вызывает уродливое KeyError. Это легко решить, добавив еще немного кода, превратив лямбда в правильную функцию.

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

Ответ 2

Здесь соответствующая ошибка/проблема: http://bugs.python.org/issue25061

Добавить встроенную поддержку enum для argparse

Я уже писал слишком много.:)

Ответ 3

Это улучшение ответа Рона Ротмана. Кроме того, переопределив __repr__ и немного изменив to_string, мы можем получить лучшее сообщение об ошибке из argparse когда пользователь вводит argparse значение.

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)

В примере Рона Ротмана, если мы передадим yellow цвет в качестве аргумента командной строки, мы получим следующую ошибку:

demo.py: error: argument color: invalid from_string value: 'yellow'

С улучшенным кодом выше, если мы передадим three в качестве аргумента командной строки, мы получим:

demo.py: error: argument some_val: invalid choice: 'three' (choose from one, two)

ИМХО, в простом случае простого преобразования имени членов перечисления в нижний регистр, метод OP кажется более простым. Однако для более сложных случаев преобразования это может быть полезно.