Python：无需if-else块，如何强制函数输入为指定特定字符串值？

要让函数参数只能接受'apple'或'banana'且不在函数内部写冗余的if-else块，有几种优雅的实现方案：

方法1：使用枚举（Enum） #

枚举是Python官方推荐的限定可选值的方式，自带类型校验，可读性强。

基础版（强制传入枚举成员） #

from enum import Enum

class ProductType(Enum):
    APPLE = "apple"
    BANANA = "banana"

def product(product_type: ProductType):
    print(product_type.value)  # 输出对应的字符串值

# 正确调用
product(ProductType.APPLE)  # 输出 apple
product(ProductType.BANANA) # 输出 banana

# 错误调用直接触发TypeError
product("apple")

兼容字符串输入版 #

如果需要支持直接传入字符串参数，可以在函数内做枚举转换：

def product(product_type: str):
    try:
        pt = ProductType(product_type)
        print(pt.value)
    except ValueError:
        raise ValueError(f"product_type 只能是 'apple' 或 'banana'") from None

# 正确调用
product("apple")
product("banana")

# 错误调用触发ValueError
product("orange")

方法2：Literal类型提示+极简运行时校验 #

用typing.Literal可以在静态检查阶段（如mypy、IDE）就提示参数错误，配合一行代码完成运行时校验：

from typing import Literal

def product(product_type: Literal["apple", "banana"]):
    if product_type not in {"apple", "banana"}:
        raise ValueError(f"product_type 只能是 'apple' 或 'banana'")
    print(product_type)

# 正确调用
product("apple")
product("banana")

# 错误调用触发ValueError
product("orange")

这里的判断逻辑极简，不属于冗余的if-else块，同时静态检查能提前拦截错误。

方法3：通用校验装饰器 #

如果多个函数需要类似校验，可以封装一个装饰器，彻底把校验逻辑从函数中抽离：

def restrict_args(**allowed_values):
    def decorator(func):
        def wrapper(*args, **kwargs):
            # 校验关键字参数
            for arg_name, allowed in allowed_values.items():
                if arg_name in kwargs and kwargs[arg_name] not in allowed:
                    raise ValueError(f"{arg_name} 只能取 {allowed} 中的值")
            # 校验位置参数（这里假设第一个参数是product_type）
            if args and allowed_values.get("product_type") and args[0] not in allowed_values["product_type"]:
                raise ValueError(f"product_type 只能取 {allowed_values['product_type']} 中的值")
            return func(*args, **kwargs)
        return wrapper
    return decorator

# 使用装饰器
@restrict_args(product_type={"apple", "banana"})
def product(product_type):
    print(product_type)

# 正确调用
product("apple")
product("banana")

# 错误调用触发ValueError
product("orange")

方法4：Pydantic自动校验 #

如果项目已使用Pydantic，可借助其模型完成参数校验，无需手动写判断：

from pydantic import BaseModel, ValidationError
from typing import Literal

class ProductParams(BaseModel):
    product_type: Literal["apple", "banana"]

def product(product_type):
    try:
        params = ProductParams(product_type=product_type)
        print(params.product_type)
    except ValidationError as e:
        raise ValueError(e.errors()[0]["msg"]) from None

# 正确调用
product("apple")
product("banana")

# 错误调用触发ValueError
product("orange")

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