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

Question

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?

Answer 1

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

Answer 2

Instead of checking for equality for each item individually, check if it's in a set of valid items.

if var in {'foo', 'bar', 'etc.'}:

Also, don't use is check for string equality. Use ==. It is possible that a string in Python can contain the same contents as another, but not be the same object. The compiler should intern (i.e. reuse) strings from literals, but this is just an optimization. Strings generated at runtime should make new objects. There are ways to intern them manually though.

---

For a very long list of options like that, I'd probably split them from a string to make the set, like

options = set("""
spam
eggs
sausage
bacon
ham
""".split())

Then you can use var in options.

You can even union other hashable types into the set.

options |= {7, 42, False}

Although user input would start out as strings anyway.

---

Another option to consider is the re module. A regular expression can match a large set of related strings, sometimes very compactly. It depends on the kinds of options you want to allow.