Marcus
Marcus

Reputation: 1023

Pydantic model_validator that can leverage the type-hints before initialisation

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

Answers (0)

Related Questions