Как заставить argparse хорошо работать с перечислениями и значениями по умолчанию? ⇐ Python

[Anonymous]

У меня есть перечисление:

Код: Выделить всё

from enum import auto, Enum

class MyEnum(Enum):
ONE = auto()
TWO = auto()
THREE = auto()

и я хочу использовать его в качестве аргумента с argparse. Чтобы быть более конкретным, я хочу создать аргумент, который принимает одно из имен перечисления («один», «два», «три») и, возможно, имеет значение по умолчанию, а соответствующий член перечисления хранится в пространстве имен. Другими словами, я хочу сделать что-то вроде этого:

Код: Выделить всё

from argparse import ArgumentParser
from enum import auto, Enum

class MyEnum(Enum):
ONE = auto()
TWO = auto()
THREE = auto()

enum_map = {e.name.lower(): e for e in MyEnum}

parser = ArgumentParser()
parser.add_argument(
"--enum",
default="two",
choices=tuple(enum_map.keys()),
help="An enum value (Default: %(default)s)",
)
args = parser.parse_args()
args.enum = enum_map[args.enum]
print(args.enum)

Это шаблонный код, и я хотел устранить его с помощью специального действия. Я вдохновился этим ответом SO и получил следующий пример:

Код: Выделить всё

from argparse import Action, ArgumentParser, FileType, Namespace
from collections.abc import Callable, Sequence
from enum import auto, Enum
from typing import Any

class EnumAction[T](Action):
_enum: type[T]
_enum_map: dict[str, T]

def __init__(
self,
option_strings: Sequence[str],
dest: str,
nargs: int | str | None = None,
const: Any = None,
default: str | T = None,
type: Callable[[str], T] | FileType | None = None,
choices: Sequence[T] | None = None,
required: bool = False,
help: str | None = None,
metavar: str | tuple[str, ...] | None = None,
) -> None:
if type is None:
raise ValueError("type must be assigned an Enum when using EnumAction")
if not issubclass(type, Enum):
raise TypeError("type must be an Enum when using EnumAction")
if choices is not None:
raise ValueError("Can't specify choices when using EnumAction")

self._enum = type
type = None
self._enum_map = {e.name.lower(): e for e in self._enum}
choices = tuple(self._enum_map.keys())
super().__init__(
option_strings,
dest,
nargs,
const,
default,
type,
choices,
required,
help,
metavar,
)

def __call__(
self,
parser: ArgumentParser,
namespace: Namespace,
values: str | Sequence[Any] | None,
option_string: str | None = None,
) -> None:
setattr(namespace, self.dest, self._enum_map[values])

class MyEnum(Enum):
ONE = auto()
TWO = auto()
THREE = auto()

parser = ArgumentParser()
parser.add_argument(
"--enum",
action=EnumAction,
type=MyEnum,
default="two",
help="An enum value (Default: %(default)s)",
)
print(parser.parse_args())

Это работает довольно хорошо, за исключением одной проблемы: если я вызываю программу без аргумента --enum, я получаю значение по умолчанию «два», но оно не обрабатывается моим действием, поэтому я получаю саму строку, а не член перечисления:

Код: Выделить всё

$ python demo.py --enum one
Namespace(enum=)

$ python demo.py
Namespace(enum='two')

Самое простое решение — установить значение по умолчанию default=MyEnum.TWO вместо «two» или заставить мое действие преобразовать параметр default в соответствующий член перечисления, если это строка. Однако это приводит к тому, что член перечисления отображается в справочном сообщении, которое не очень читабельно (по этой причине я задаю имена членов вместо реальных членов):

Код: Выделить всё

$ python demo.py -h
usage: demo.py [-h] [--enum {one,two,three}]

options:
-h, --help            show this help message and exit
--enum {one,two,three}
An enum value (Default: MyEnum.TWO)

Как я могу получить член перечисления в пространстве имен как при передаче аргумента, так и при использовании значения по умолчанию, и видеть имя члена в справочном сообщении вместо самого члена?


Подробнее здесь: https://stackoverflow.com/questions/787 ... ult-values