Pydantic多模型基于Enum实现不同字段选项的优化方案咨询

问题背景 #

使用Pydantic（v2.6.1）定义模型时，若需对同一个枚举类型的字段，在不同模型中设置不同的可选值范围，传统的field_validator方式会产生大量冗余代码。例如：

from enum import Enum
from pydantic import BaseModel, field_validator

class FruitEnum(str, Enum):
    pear = 'pear'
    banana = 'banana'
    grapes = 'grapes'

class CookingModel_1(BaseModel):
    fruit: FruitEnum

    @field_validator('fruit')
    def _validate_fruit(cls, value):
        if value not in (FruitEnum.pear, FruitEnum.grapes):
            raise ValueError(f'{value} is not an allowed fruit!')

class CookingModel_2(BaseModel):
    fruit: FruitEnum

    @field_validator('fruit')
    def _validate_fruit(cls, value):
        if value not in (FruitEnum.banana, FruitEnum.grapes):
            raise ValueError(f'{value} is not an allowed fruit!')

即使通过嵌套函数生成验证器减少冗余，当涉及多字段、多枚举类型时，代码仍不够简洁：

def get_fruit_validator(allowed_fruits):
    def _validator(fruit):
        if fruit not in allowed_fruits:
            raise ValueError(f'{fruit} is not an allowed fruit!')
        return fruit
    return _validator

class CookingModel_1(BaseModel):
    fruit: FruitEnum
    _fruit_validator = field_validator('fruit')(get_fruit_validator({FruitEnum.pear, FruitEnum.grapes}))

class CookingModel_2(BaseModel):
    fruit: FruitEnum
    _fruit_validator = field_validator('fruit')(get_fruit_validator({FruitEnum.banana, FruitEnum.grapes}))

以下是几种无需自定义验证器的优雅解决方案：

方案1：使用Literal直接指定允许的枚举成员 #

Pydantic支持将Literal与枚举成员结合，直接限制字段的可选值，无需额外编写验证器，Pydantic会自动完成验证并返回友好的错误信息：

from enum import Enum
from pydantic import BaseModel, Literal

class FruitEnum(str, Enum):
    pear = 'pear'
    banana = 'banana'
    grapes = 'grapes'

class CookingModel_1(BaseModel):
    fruit: Literal[FruitEnum.pear, FruitEnum.grapes]

class CookingModel_2(BaseModel):
    fruit: Literal[FruitEnum.banana, FruitEnum.grapes]

优点：代码最简洁，无需额外逻辑，Pydantic原生支持，错误提示清晰。
缺点：若同一枚举子集需在多个模型中复用，会存在重复代码。

方案2：动态生成受限枚举子类 #

如果需要复用特定的枚举值子集，可以编写一个工具函数，动态生成继承自原枚举的子类，仅包含指定的允许成员。这样每个模型只需使用对应的受限枚举类型即可：

from enum import Enum, EnumMeta
from pydantic import BaseModel

class FruitEnum(str, Enum):
    pear = 'pear'
    banana = 'banana'
    grapes = 'grapes'

def create_restricted_enum(base_enum: EnumMeta, allowed_members: list[Enum]) -> EnumMeta:
    # 构建受限枚举的成员字典
    restricted_members = {member.name: member.value for member in allowed_members}
    # 动态生成枚举子类
    return Enum(f"Restricted{base_enum.__name__}", restricted_members, type=base_enum)

# 生成受限枚举
PearGrapesEnum = create_restricted_enum(FruitEnum, [FruitEnum.pear, FruitEnum.grapes])
BananaGrapesEnum = create_restricted_enum(FruitEnum, [FruitEnum.banana, FruitEnum.grapes])

class CookingModel_1(BaseModel):
    fruit: PearGrapesEnum

class CookingModel_2(BaseModel):
    fruit: BananaGrapesEnum

优点：枚举子集可复用，模型定义更清晰，符合枚举类型的语义。
缺点：需要额外编写工具函数，适合有复用需求的场景。

方案3：使用Annotated结合BeforeValidator（可选） #

如果需要更灵活的验证逻辑（比如动态调整允许值），可以结合Annotated和BeforeValidator，配合参数化的验证函数，相比传统field_validator更简洁：

from enum import Enum
from pydantic import BaseModel, BeforeValidator, Annotated

class FruitEnum(str, Enum):
    pear = 'pear'
    banana = 'banana'
    grapes = 'grapes'

def validate_allowed_fruits(allowed_fruits):
    def validator(value):
        if value not in allowed_fruits:
            raise ValueError(f'{value} is not an allowed fruit!')
        return value
    return validator

# 用Annotated封装字段类型和验证器
FruitPearGrapes = Annotated[FruitEnum, BeforeValidator(validate_allowed_fruits({FruitEnum.pear, FruitEnum.grapes}))]
FruitBananaGrapes = Annotated[FruitEnum, BeforeValidator(validate_allowed_fruits({FruitEnum.banana, FruitEnum.grapes}))]

class CookingModel_1(BaseModel):
    fruit: FruitPearGrapes

class CookingModel_2(BaseModel):
    fruit: FruitBananaGrapes

优点：验证逻辑可复用，字段类型封装更直观，适合复杂验证场景。
缺点：仍需编写基础验证函数，但相比每个模型定义验证器更简洁。

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