使枚举在 pydantic 中更灵活

Question

在以下模型中

from pydantic import (BaseModel, validator)
from enum import Enum

class City(str, Enum):
    new_york = "New York"
    los_angeles = "Los Angeles"

class CityData(BaseModel):
    city:City
    population:int

可以将 CityData 的实例构造为

c = CityData(city="New York", population=8419000)

我希望能够创建与

相同的实例

c = CityData(city="NY", population=8419000)

问题：如何做到这一点？

向 CityData 添加验证器，如

@validator("city")
def _flexible_city(cls, v, values, *kwargs):
    if v == "NY":
        return "New York"
    else:
        return v

不起作用。看起来值 v 已经是 City 的一个实例。这是否意味着转换应该发生在 City 的构造函数中？

Answer 1

我想我找到了使用 City

的 _missing_ 方法的方法


class City(str,Enum):
    new_york = "New York"
    los_angeles = "Los Angeles"
    @classmethod
    def _missing_(cls, name):
        if name == "NY":
            return cls.new_york
        else:
            return cls[name]

Answer 2

在这种情况下，使用在解析和其他验证之前调用的 pre-validator 会很方便。

class CityData(BaseModel):
    city: City
    population: int

    @validator("city", pre=True)
    def _flexible_city(cls, v):
        if v == "NY":
            return "New York"
        else:
            return v

使枚举在 pydantic 中更灵活

Making an Enum more flexible in pydantic

python

pydantic