pythonparameter-passingargparsepredefined-variables

How do I restrict passed variable to specific values in python?


I am writing the script where I pass values with CLI through argparse module. I am wondering if this is possible to restrict variable to hold pre-defined values, to avoid user mistake. It is not the type restriction, values are consists of letters as well as digits, surely I can write an if block, but I have about 30 pre-defined values, so writing something like

if var is value1 or var is value2 ... or var is value30:
  pass
else:
    print("oops, your value does not fit")

would be painful. What is the proper way of doing this?


Solution

  • With choices:

    In [214]: parser = argparse.ArgumentParser()
    In [215]: parser.add_argument('--foo', choices=['one','two','three','four']);
    

    Accepted:

    In [216]: parser.parse_args('--foo one'.split())
    Out[216]: Namespace(foo='one')
    

    rejected:

    In [217]: parser.parse_args('--foo five'.split())
    usage: ipython3 [-h] [--foo {one,two,three,four}]
    ipython3: error: argument --foo: invalid choice: 'five' (choose from 'one', 'two', 'three', 'four')
    

    help:

    In [218]: parser.parse_args('-h'.split())
    usage: ipython3 [-h] [--foo {one,two,three,four}]
    
    optional arguments:
      -h, --help            show this help message and exit
      --foo {one,two,three,four}
    

    If I'd defined a metavar, the help will be

    usage: ipython3 [-h] [--foo CHOICES]
    
    optional arguments:
      -h, --help     show this help message and exit
      --foo CHOICES
    

    Or if the choices is too long, define a type function:

    In [222]: def mychoices(astr):
         ...:     if astr in ['one','two','three','four']:
         ...:         return astr
         ...:     else:
         ...:         raise argparse.ArgumentTypeError('Wrong choice')
    
    In [223]: parser = argparse.ArgumentParser()
    In [224]: parser.add_argument('--foo', type=mychoices);
    
    In [225]: parser.parse_args('--foo one'.split())
    Out[225]: Namespace(foo='one')
    
    In [226]: parser.parse_args('--foo five'.split())
    usage: ipython3 [-h] [--foo FOO]
    ipython3: error: argument --foo: Wrong choice
    
    In [227]: parser.parse_args('-h'.split())
    usage: ipython3 [-h] [--foo FOO]
    
    optional arguments:
      -h, --help  show this help message and exit
      --foo FOO