Reputation: 1023
I have a use-case where I will use Pydantic classes to send information between different services using a messaging service (PubSub). Since the pydantic classes as serializable, Pydantic's BaseModel becomes very convinient to use to serialize and deserialize data. However, as part of my models I have a lot of Enums, which I want to support to send as well.
Below is a minimalistic example, but I have yet to find a solution that works. I'm hoping someone has more pydantic & typing knowledge than me :) Python 3.11 & Pydantic 2.8 used.
from pydantic import BaseModel, Field, model_validator
from typing import Annotated, Any, Type, Optional
from enum import Enum
def transform_str_to_enum(value: str, enum_type: Type[Enum]) -> Enum:
"""Transform a string to an Enum. Raise a KeyError error if the value is not a valid Enum member."""
return enum_type[value.upper()]
class MyEnum(Enum):
A = "option1"
B = "option2"
class MyClass(BaseModel):
str_field: Annotated[str, Field(description="A normal string field")]
int_field: Annotated[int, Field(description="A normal int field")]
# Cannot change type-hinting to MyEnum | str, and use @model_validator(mode="after"),
# because then the type hinting will be wrong for later use that depend on this field being an Enum.
enum_field: Annotated[
MyEnum,
Field(description="MyEnum, can be either A or B, and be sent as both string or Enum"),
]
optional_enum_field: Annotated[
Optional[MyEnum],
Field(description="Optional MyEnum, can be either A or B or None, and be sent as both string or Enum")
]
# @model_validator(mode="before")
# def validate_enums(cls, values: dict[str, Any]) -> dict[str, Any]:
# """
# Validate all Enums here. If the value is type-hinted as a subclass of Enum, but the value is a string, we convert it to the Enum.
# """
# # Pesudo code, doesnt work. Implement this in a way that works.
# # for field, value in values.items():
# # if field is type-hinted as a class that is a subclass of Enum:
# # values[field] = transform_str_to_enum(value=value, enum_type=field_class)
# return values
# Test that it works
expected_instance = MyClass(
str_field="test",
int_field=1,
enum_field=MyEnum.A,
optional_enum_field=MyEnum.B
)
def test_that_it_works_with_enum():
created_instance = MyClass(
str_field="test",
int_field=1,
enum_field=MyEnum.A,
optional_enum_field=MyEnum.B
)
assert created_instance == expected_instance
def test_that_it_works_with_str():
created_instance = MyClass(
str_field="test",
int_field=1,
enum_field="A",
optional_enum_field="B"
)
assert created_instance == expected_instance
# Both should pass.
test_that_it_works_with_enum()
test_that_it_works_with_str()
Any clever ideas?
Upvotes: 1
Views: 129