Making an Enum more flexible in pydantic

Question

In the following model

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

One can construct instances of CityData as

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

I would like to be able to create the same instance with

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

Question: How to do this?

Adding a validator to CityData, like

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

doesn't work. It looks like the value v is already an instance of City. Does that mean that the conversion should happen in the constructor of City?

Answer 1

In such cases, it is convenient to use a pre-validator that is called prior to parsing and other validations.

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