Pydantic field validation against Enum name

Question

I have an Enum:

class DataType(Enum):
    TIMESTAMP = MyClass("ts")
    DATETIME = MyClass("dt")
    DATE = MyClass("date")

and a Pydantic Model with a field of the type of that Enum:

class Aggregation(BaseModel):
    field_data_type: DataType

Is there a convenient way to tell Pydantic to validate the Model field against the names of the enum rather than the values? i.e., to be able to build this Model:

agg = Aggregation(field_data_type="TIMESTAMP")

Answer 1

You can use the built-in functionality of Enum to accomplish this:

agg = Aggregation(field_data_type=DataType["TIMESTAMP"])

You can wrap this in a function and then create a new annotated type that calls it:

def convert_to_data_type(s: str):
    return DataType[s]

StrDataType = Annotated[str, AfterValidator(convert_to_data_type)]

class Aggregation(BaseModel):
    field_data_type: StrDataType

agg = Aggregation(field_data_type="TIMESTAMP")