Pythonic way to have a choice of 2-3 options as an argument to a function

robintw picture robintw · Mar 1, 2012 · Viewed 17.3k times · Source

I have a Python function which requires a number of parameters, one of which is the type of simulation to perform. For example, the options could be "solar", "view" or "both.

What is a Pythonic way to allow the user to set these?

I can see various options:

  1. Use a string variable and check it - so it would be func(a, b, c, type='solar')

  2. Set some constants in the class and use func(a, b, c, type=classname.SOLAR)

  3. If there are only two options (as there are for some of my functions) force it into a True/False argument, by using something like func(a, b, c, do_solar=False) to get it to use the 'view' option.

Any preferences (or other ideas) for Pythonic ways of doing this?

Answer

Steven T. Snyder picture Steven T. Snyder · Mar 1, 2012

If the point Niklas' makes in his answer doesn't hold, I would use a string argument. There are Python modules in the standard library that use similar arguments. For example csv.reader().

sim_func(a, b, c, sim_type='solar')

Remember to give a reasonable error inside the function, that helps people out if they type in the wrong thing.

def sim_func(a, b, c, sim_type='solar'):
    sim_types = ['solar', 'view', 'both']
    if sim_type not in sim_types:
        raise ValueError("Invalid sim type. Expected one of: %s" % sim_types)
    ...