it-swarm.com.de

Unterstützung für Aufzählungsargumente in argparse

Gibt es eine bessere Möglichkeit, Enums als Argumenttypen zu unterstützen als dieses Muster?

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

Ich sehe, das ist eine alte Frage, aber ich bin gerade auf das gleiche Problem gestoßen (Python 2.7) und hier habe ich es gelöst:

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

Beachten Sie, dass das Definieren von __str__ erforderlich ist, um die Hilfeausgabe von ArgumentParser zu erhalten, die die lesbaren Werte (Color) des Benutzers enthält.

Einige Beispielaufrufe:

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

Da die Frage des OP Ganzzahlen als Werte angibt, ist hier eine leicht modifizierte Version, die in diesem Fall funktioniert (unter Verwendung der Aufzählungsnamen anstelle der Werte als Befehlszeilenargs):

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

Der einzige Nachteil besteht darin, dass ein fehlerhafter Parameter eine unschöne KeyError verursacht. Das lässt sich leicht lösen, indem nur etwas mehr Code hinzugefügt wird, um das Lambda in eine richtige Funktion umzuwandeln.

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

Hier ist der relevante Fehler: http://bugs.python.org/issue25061

Native Enum-Unterstützung für argparse hinzufügen

Ich habe dort schon zu viel geschrieben. :) 

0
hpaulj

Dies ist eine Verbesserung gegenüber Ron Rothmans Antwort . Wenn Sie auch __repr__ überschreiben und to_string ein wenig ändern, erhalten Sie eine bessere Fehlermeldung von argparse, wenn der Benutzer einen falschen Wert eingibt.

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)

Wenn wir in Ron Rothmans Beispiel die Farbe yellow als Befehlszeilenargument übergeben, erhalten wir den folgenden Fehler:

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

Mit dem oben beschriebenen verbesserten Code erhalten wir, wenn wir three als Befehlszeilenargument übergeben:

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

IMHO, im einfachen Fall, dass nur der Name der Enum-Mitglieder in Kleinbuchstaben umgewandelt wird, scheint die OP-Methode einfacher zu sein. Bei komplexeren Konvertierungsfällen kann dies jedoch hilfreich sein.

0
David Lechner