Pythonic方式可以选择2-3个选项作为函数的参数

Question

我有一个Python函数,需要许多参数,其中一个参数是要执行的模拟类型.例如,选项可以是"太阳能","视图"或"两者".

什么是Pythonic方式允许用户设置这些？

我可以看到各种选择:

1. 使用字符串变量并检查它 - 所以它会 func(a, b, c, type='solar')

2. 在类中设置一些常量并使用 func(a, b, c, type=classname.SOLAR)

3. 如果只有两个选项(就像我的一些函数一样)强制它进入一个True/False参数,通过使用类似的东西func(a, b, c, do_solar=False)来使它使用'view'选项.

Pythonic做这个的任何偏好(或其他想法)？

Answer 1

我推荐在Python 3.8typing.Literal中实现的内置

简单类型提示

第一个示例仅在 IDE 中进行“类型”提示

from typing import Literal

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

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

这是它在 PyCharm 中的样子，注意“solra”是如何突出显示的以及它如何很好地为我们生成文档

简单断言

如果你想在参数检查失败时引发异常，我们可以利用get_args保持干燥

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}"

动态断言

更进一步，我们可以创建一个函数来动态检查任何提供的文字

我们可以从中获取前一帧sys以提取给定函数中提供的参数。然后我们使用函数的__annotations__属性来查看参数应该是什么值

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'

2022 年 10 月编辑：从inspect.stack()[1](963 usec) 更改为sys._getframe(1)(1.2 usec)。应该inspect.currentframe().f_back从一开始就使用（2.7 usec），但sys更好！

2023 年 2 月编辑：修复了定义返回类型时的错误，由@jung rhew指出。还发布了通用工具以使其可维护。

Answer 2

如果Niklas'在他的回答中提出的观点不成立,我会使用一个字符串参数.标准库中有Python模块使用类似的参数.例如csv.reader().

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

记住在函数内部给出一个合理的错误,如果他们键入错误的东西,可以帮助人们.

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 3

我不喜欢任何这些选择.

我定义了两个不同的函数,perform_solar(a, b, c)并perform_view(a, b, c)让调用者决定他想要使用哪些函数,以及使用哪些参数.

如果您认为必须将这些函数打包到一个函数中的原因是它们共享状态,则应该在对象中共享该状态并将函数定义为方法.

Answer 4

您可以像这样使用断言语句：

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

如果 sim_types 不在列表中，python 将引发断言错误

Answer 5

只是根据您的选项#1（“使用字符串变量并检查它”）为此编写了一个装饰器工厂。

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

所以现在我们可以创建这样的新函数：

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

并测试一下：

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

输出：

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']'

这个装饰器仍然需要改进，但现在已经可以使用了。

这里进行了改进的修订：

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