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 \"bot

Answer 1

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