Demo of creating custom argument types for argparse (percent and Enum)

argparse_demo.py

#
# argparse_demo.py
#
# Paul McGuire - April, 2019
#
# Presented as a lightning talk at PyTexas 2019
#
# Gist url: https://gist.github.com/ptmcg
#
import sys
import argparse

parser = argparse.ArgumentParser()

# create a parser that will accept arguments:
# - name
# - qty
# - pct
# - color

# define a str argument
parser.add_argument('--name')

#~ parser.add_argument('--qty')
#~ parser.add_argument('--pct')
#~ parser.add_argument('--color')

# define an argument of a built-in type
parser.add_argument('--qty', type=int)

# define an argument of a custom type
# - create a method that takes a str and returns
#   the desired type
def percent_type(s: str) -> float:
    try:
        value = float(s)
    except TypeError:
        raise argparse.ArgumentTypeError(
            f"{s!r} is not a valid numeric value")

    if not (0 <= value <= 100):
        raise argparse.ArgumentTypeError(
            f"{s} must be in the range 0-100")

    return value / 100

parser.add_argument('--pct', type=percent_type)


# define an arg that is an enumerated value

from enum import Enum

class ArgTypeMixin(Enum):

    @classmethod
    def argtype(cls, s: str) -> Enum:
        try:
            return cls[s]
        except KeyError:
            raise argparse.ArgumentTypeError(
                f"{s!r} is not a valid {cls.__name__}")

    def __str__(self):
        return self.name


class Color(ArgTypeMixin, Enum):
    red = 0
    orange = 1
    yellow = 2
    green = 3
    blue = 4
    purple = 5
    gold = 6


parser.add_argument('--color', 
                    type=Color.argtype, 
                    choices=Color)


if __name__ == "__main__":

    # use argparse parser to parse input args
    if not sys.argv[1:]:
        sys.exit(0)
    parsed_args = parser.parse_args(sys.argv[1:])


    # show the values returned from parsing, already converted
    # to the correct types
    for k, v in vars(parsed_args).items():
        print(f"{k:>8}: {repr(v):12} {type(v)}")


    # use parsed values without further conversion
    if parsed_args.qty is not None and parsed_args.qty > 100:
        print(f"{parsed_args.qty}?! That's a lot!")

    if parsed_args.qty is not None and parsed_args.pct is not None:
        print("That comes out to "
              f"{parsed_args.pct * parsed_args.qty}")

    if parsed_args.color is Color.gold:
        print("Pretty!")

Thanks, I was trying to find the best way to use an Enum with argparse and I was having problems.

I ended up using a simpler approach by using a string-type enum, and simply overriding __str__ to give the value:

class TLSMode(Enum):
    NONE = 'none'
    STARTTLS = 'starttls'
    ALWAYS = 'always'

    def __str__(self):
        return self.value

def parse_args():
    ap = ArgumentParser()
    ap.add_argument('--tls', choices=TLSMode, type=TLSMode)
    return ap.parse_args()

Very nice.

Thanks, I was trying to find the best way to use an Enum with argparse and I was having problems.

I ended up using a simpler approach by using a string-type enum, and simply overriding __str__ to give the value:

class TLSMode(Enum):
    NONE = 'none'
    STARTTLS = 'starttls'
    ALWAYS = 'always'

    def __str__(self):
        return self.value

def parse_args():
    ap = ArgumentParser()
    ap.add_argument('--tls', choices=TLSMode, type=TLSMode)
    return ap.parse_args()

Excelent