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

Question

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 1

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

Answer 2

I recommend built-in typing.Literal which was implemented in Python 3.8

Simple type hinting

This first example only does "type" hinting in the IDE

from typing import Literal

_TYPES = Literal["solar", "view", "both"]

def func(a, b, c, type_: _TYPES = "solar"):
    pass

Here's how it looks in PyCharm, notice how "solra" is highlighted and how it nicely generates the docs for us

Simple assertion

If you want to raise an exception when the parameter check fails we can make use of get_args to stay dry

from typing import Literal, get_args

_TYPES = Literal["solar", "view", "both"]

def func(a, b, c, type_: _TYPES = "solar"):
    options = get_args(_TYPES)
    assert type_ in options, f"'{type_}' is not in {options}"

Dynamic assertion

Taking it one step further we can create a function to dynamically check any supplied Literals

We can get the previous frame from sys to extract the provided arguments in the given function. We then use the function's __annotations__ attribute to see what values the arguments should be

from typing import Literal, get_args, get_origin
from sys import _getframe

def enforce_literals(function):
    kwargs = _getframe(1).f_locals
    for name, type_ in function.__annotations__.items():
        value = kwargs.get(name)
        options = get_args(type_)
        if get_origin(type_) is Literal and name in kwargs and value not in options:
            raise AssertionError(f"'{value}' is not in {options} for '{name}'")

_TYPES = Literal["solar", "view", "both"]
_NUMS = Literal[1, 2, 3, 4, 5]

def func(a, b, c, type_: _TYPES = "solar", num: _NUMS = 1):
    enforce_literals(func)

func(1, 2, 3, "solar", 6)

AssertionError: '6' is not in (1, 2, 3, 4, 5) for 'num'

Oct 2022 Edit: Changed from inspect.stack()[1] (963 usec) to sys._getframe(1) (1.2 usec). Should have used inspect.currentframe().f_back (2.7 usec) from the start but sys is even better!

Feb 2023 Edit: Fixed bug for when return type is defined, pointed out by @jung rhew. Also published generaltool to make it maintainable.

Answer 3

I don't like any of those options.

I'd define two different functions, perform_solar(a, b, c) and perform_view(a, b, c) and let the caller decide which ones he wants to use, in which order and with which arguments.

If the reason why you thought you'd have to pack these into one single function is that they share state, you should share that state in an object and define the functions as methods.

Answer 4

You can use the assert statement like this:

assert sim_types in ['solar', 'view', 'both'], 'sim type parameter must be solar, view or both'

If sim_types is not in the list, python will raise an Assertion Error

Answer 5

Just have written a decorator factory for this, based on your option #1, which is "Use a string variable and check it".

def limited_argument_choices(choices: Dict[int or str, Iterable] = None) -> Callable:
    """decorator factory: force arguments of a func limited in the given choices

    :param choices: a dict which describes the choices for the value-limited arguments.
            the key of the dict must be either the index of args or the key_str of kwargs,
            while the value of the dict must be an iterable."""
    err_fmt = "value of '{}' is not a valid choice: '{}'"

    def decorator(func):
        if not choices:
            return func

        @wraps(func)
        def decorated_func(*args, **kwargs):
            for i in range(len(args)):
                if i in choices and args[i] not in choices[i]:
                    param_name = func.__code__.co_varnames[i]
                    valid_choices = list(choices[i])
                    raise ValueError(err_fmt.format(param_name, valid_choices))
            for k in kwargs:
                if k in choices and kwargs[k] not in choices[k]:
                    raise ValueError(err_fmt.format(k, list(choices[k])))

            return func(*args, **kwargs)

        return decorated_func

    return decorator

So now we could make new functions like this:

@limited_argument_choices({1: (0, 1, 2), 'y': ('hi', 'hello')})
def test(a, b, c, y=1):
    print(a, b, c, y)

And test it:

test(0, 1, 2, y='hello')
test(0, 3, 2, y='hello')
test(0, 1, 2, y='world')

Output:

0 1 2 hello
ValueError: value of 'b' is not a valid choice: '[0, 1, 2]'
ValueError: value of 'y' is not a valid choice: '['hi', 'hello']'

This decorator still needs improvements, but it's already usable now.

---

An improved revision here:

def decorator_factory_args_choices(choices: Dict[int or str, Iterable]) -> Decorator:
    """decorator factory: force arguments of a func limited inside the given choices

    :param choices: a dict which describes the choices of arguments
        the key of the dict must be either the index of args or the key(str) of kwargs
        the value of the dict must be an iterable."""
    err_fmt = "value of '{}' is not a valid choice in {}"

    def decorator(func):
        @wraps(func)
        def decorated_func(*args, **kwargs):
            for arg_index in range(len(args)):
                param_name = func.__code__.co_varnames[arg_index]
                if arg_index in choices and args[arg_index] not in choices[arg_index]:
                    raise ValueError(err_fmt.format(param_name, choices[arg_index]))
                elif param_name in choices and args[arg_index] not in choices[param_name]:
                    raise ValueError(err_fmt.format(param_name, choices[param_name]))
            for param_name in kwargs:
                if param_name in choices and kwargs[param_name] not in choices[param_name]:
                    raise ValueError(err_fmt.format(param_name, choices[param_name]))

            return func(*args, **kwargs)

        return decorated_func

    return decorator

Answer 6

You can use optional (keyword) arguments like this

def func(a, b, c, **kw):
    if kw.get('do_solar'):
        # Do Solar
    if kw.get('do_view'):
        # Do view

Answer 7

Since functions are objects in python, you could actually process *args as a list of methods and pass the types of simulations as arbitratry args at the end. This would have the benefit of allowing you to define new simulations in the future without having to refactor this code.

def func(a, b, c, *args):
    for arg in args:
        arg(a, b, c)

def foosim(a, b, c):
    print 'foosim %d' % (a + b + c)

def barsim(a, b, c):
    print 'barsim %d' % (a * b * c)

Use:

func(2, 2, 3, foosim)
func(2, 2, 3, barsim)
func(2, 2, 3, foosim, barsim)

Output:

foosim 7
barsim 12
foosim 7
barsim 12

Answer 8

I really like dictionaries (the pythonic switch-case replacement) for that because it is performant, easy to read, and easy to maintain/extend:

def example_function(a, b, c, op='add'):
    return {'add': a+b+c, 'multiply': a*b*c}[op]


example_function(0, 1, 2, 'add') # returns 3
example_function(0, 1, 2, 'multiply') # returns 0

And it fails nicely for wrong/not implemented type-parameters:

example_function(0, 1, 2, 'divide') # returns a key error for 'divide'

Answer 9

The answer of @Mandera is pretty much what you need to use today, if you use Python >= 3.8

But I think the answer is not yet use-ready. Thus, I have converted it to a decorator, which executes the function after all the Literal checks have been passed.

Since I turned it into a decorator, the originally used function _getframe was not working as intended. I've replaced it with getfullargspec. Still no need for third-party packages.

Wanted to comment to the Mandera's answer, but have no reputation on the new account, duh. Had to post it as a separate answer.

from typing import Literal, get_args, get_origin
from inspect import getfullargspec


def enforce_literals(function):
    """Decorator that raises AssertionError on Literal check failure."""

    def decorator(*args, **kwargs):
        specs = getfullargspec(function)
        args = {key: args[i] for i, key in enumerate(specs.args) if i < len(args)}
        #  key_values = args | kwargs  # use this if python >= 3.9
        key_values = {**args_new, **kwargs}  # this is for python 3.8

        for name, type_ in getfullargspec(function).annotations.items():
            value = key_values[name]
            options = get_args(type_)
            if (
                get_origin(type_) is Literal
                and name in specs.args
                and value not in options
            ):
                raise AssertionError(f"'{value}' is not in {options} for '{name}'")
        return function(*args, **kwargs)

    return decorator


_TYPES = Literal["solar", "view", "both"]
_NUMS = Literal[1, 2, 3, 4, 5]


@enforce_literals
def func(a, b, c, type_: _TYPES = "solar", num: _NUMS = 5):
    print("Function executed")

func(1, 2, 3, "solar", num=5)
func(1, 2, 3, "solar", 6)