Pydantic中的枚举选择

Question

我有一个使用 pydantic（版本 2.6.1）创建的类，这个类里有一个属性，我想限制用户可以选择的选项。我可以用枚举（enum）来实现这个目的。但问题是，我有多个类需要设置不同的选择。下面是一个例子：

from enum import Enum
from pydantic import BaseModel, ValidationError, 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 a 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 a allowed fruit.!')

对于我的 CookingModel_1 和 CookingModel_2 类，fruit 的选择是不同的。

我可以像上面那样来完成这个任务。但是如果我的烹饪模型列表很长（比如有 15 个），我就得在每个类里都定义一个验证器，这样会产生很多重复的代码，做的几乎是同样的事情。

基于这个想法，我可以定义一个嵌套函数，来创建一个参数化的字段验证器，像下面这样，从而减少重复的代码。

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


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


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

这对我的使用场景来说已经足够好了。

不过，如果我有多个字段需要其他枚举和长列表，这样还是会导致代码重复。我在寻找一种更“优雅”的解决方案。也许可以引入一种新的约束类型（ConstraintType），来帮助限制选择？有没有其他不需要编写自定义验证器的替代方案呢？

Answer 1

感谢@yaakov-bressler的帮助，我用typing_extensions里的Literal解决了这个问题，下面是代码。

from typing_extensions import Literal

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

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