如何以Python风格限制函数输入参数的取值范围？

除了直接用assert或if in做硬编码校验外，自定义可复用装饰器是最符合Python风格的实现方式——它能把校验逻辑和业务逻辑解耦，还能在多个函数间复用校验规则。

实现通用参数校验装饰器 #

借助functools.wraps保留原函数的元信息，结合inspect模块解析函数参数，我们可以写一个通用的校验装饰器：

from functools import wraps
import inspect

def validate_param(param_name, allowed_values):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            # 绑定函数参数，解析传入的参数值
            sig = inspect.signature(func)
            bound_args = sig.bind(*args, **kwargs)
            bound_args.apply_defaults()
            
            # 获取目标参数的值
            try:
                param_value = bound_args.arguments[param_name]
            except KeyError:
                raise ValueError(f"函数未定义参数 {param_name}")
            
            # 校验取值范围
            if param_value not in allowed_values:
                raise ValueError(
                    f"参数 {param_name} 取值非法：必须是 {allowed_values} 中的一个，当前值为 {param_value}"
                )
            return func(*args, **kwargs)
        return wrapper
    return decorator

装饰目标函数 #

直接给函数加上装饰器，指定要校验的参数名和允许的取值范围即可：

@validate_param("a", {1, 2, 3})
def f(a, b):
    # 这里写你的业务逻辑
    return a * b

效果测试 #

调用时如果传入不符合要求的参数，会直接抛出明确的异常：

f(4, 10)  # 抛出 ValueError: 参数 a 取值非法：必须是 {1, 2, 3} 中的一个，当前值为 4

额外方案：使用第三方库简化实现 #

如果项目中已经引入了pydantic这类数据校验库，也可以用它的函数装饰器快速实现，代码会更简洁：

from pydantic import validate_arguments

@validate_arguments
def f(a: int, b):
    if a not in {1,2,3}:
        raise ValueError("参数a必须在[1,2,3]范围内")
    # 业务逻辑
    return a + b

这种方式借助了第三方库的参数解析能力，也是Python生态中常用的实践。

内容的提问来源于stack exchange，提问作者f1msch