web-dev-qa-db-fra.com

Prise en charge des arguments Enum dans argparse

Existe-t-il un meilleur moyen de supporter Enums en tant que types d’arguments argparse que ce modèle?

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()]
pythonargparse
18
Andrzej Pronobis

Je vois que c'est une vieille question, mais je viens de rencontrer le même problème (Python 2.7) et voici comment je l'ai résolu:

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

Notez que la définition de __str__ est requise pour que la sortie d'aide de ArgumentParser inclue les valeurs lisibles par l'homme de Color.

Quelques exemples d'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}

Comme la question de l'OP spécifiait des entiers sous forme de valeurs, voici une version légèrement modifiée qui fonctionne dans ce cas (en utilisant les noms enum plutôt que les valeurs, comme l'argument de la ligne de commande):

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

Le seul inconvénient est qu’un paramètre incorrect provoque une KeyError moche. Cela est facilement résolu en ajoutant juste un peu plus de code, en convertissant le lambda en une fonction appropriée.

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))
37
ron rothman

Voici le bug/problème correspondant: http://bugs.python.org/issue25061

Ajouter un support enum natif pour argparse

J'ai déjà trop écrit là-bas. :) 

0
hpaulj

Ceci dans une amélioration sur réponse de Ron Rothman . En remplaçant également __repr__ et en modifiant to_string un peu, nous pouvons obtenir un meilleur message d'erreur de argparse lorsque l'utilisateur entre une valeur incorrecte.

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)

Dans l'exemple de Ron Rothman, si nous passons la couleur yellow comme argument de ligne de commande, nous obtenons l'erreur suivante:

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

Avec le code amélioré ci-dessus, si nous passons three comme argument de ligne de commande, nous obtenons:

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

IMHO, dans le cas simple de convertir simplement le nom des membres de l'énum en minuscule, la méthode du PO semble plus simple. Toutefois, pour les cas de conversion plus complexes, cela pourrait être utile.

0
David Lechner