Using Pydantic models within enums
Posted on 2025-04-26 in Programmation
In previous articles, I provided some tips on enums and explained how to use a custom class with Pydantic. Now, I’ll dig into all kinds of enum usages with Pydantic, including enums whose members are Pydantic models themselves! Let’s dive right in!
Just like before, you can access the code of this article as a Python script and as a notebook. I won’t include the imports in the code examples for brevity, refer to the notebook or the script if needed.
Using normal enums
Let’s start small with how Pydantic behaves with normal enums.
class IntValues(IntEnum): first = 1 second = 2 class StrValues(StrEnum): valid = "valid" invalid = "invalid" class NormalEnumModel(BaseModel): # False is the default. If True, will use 1 or "valid" # in the model instead of the enum members. model_config = ConfigDict(use_enum_values=False) int_value: IntValues str_value: StrValues m = NormalEnumModel.model_validate({"int_value": 1, "str_value": "valid"}) print("The model:", m) print("Dumped in Python mode:", m.model_dump(mode="python")) print("Dumped in JSON mode:", m.model_dump(mode="json"))
The model: int_value=<IntValues.first: 1> str_value=<StrValues.valid: 'valid'> Dumped in Python mode: {'int_value': <IntValues.first: 1>, 'str_value': <StrValues.valid: 'valid'>} Dumped in JSON mode: {'int_value': 1, 'str_value': 'valid'}
Each field is correctly parsed as an enum member. When dumping in JSON mode, the value is used and the member remains in Python mode. I say this behaves as expected.
Once again, as expected, values outside the enum will raise a ValidationError:
try: # Values outside the enum, will fail! NormalEnumModel.model_validate({"int_value": 10, "str_value": "nope"}) assert False, "Must fail" except ValidationError as e: print("Failing as expected with values outside the enum with:") print(e)
Failing as expected with values outside the enum with: 2 validation errors for NormalEnumModel int_value Input should be 1 or 2 [type=enum, input_value=10, input_type=int] For further information visit https://errors.pydantic.dev/2.11/v/enum str_value Input should be 'valid' or 'invalid' [type=enum, input_value='nope', input_type=str] For further information visit https://errors.pydantic.dev/2.11/v/enum
Let’s illustrate the behavior when forcing Pydantic to use enum values instead of enum members with use_enum_values=True:
class MyModelWithValues(BaseModel): model_config = ConfigDict(use_enum_values=True) int_value: IntValues str_value: StrValues m = MyModelWithValues.model_validate({"int_value": 1, "str_value": "valid"}) print("The model:", m) print("Dumped in Python mode:", m.model_dump(mode="python")) print("Dumped in JSON mode:", m.model_dump(mode="json"))
The model: int_value=1 str_value='valid' Dumped in Python mode: {'int_value': 1, 'str_value': 'valid'} Dumped in JSON mode: {'int_value': 1, 'str_value': 'valid'}
This time, the values is used in the model and no matter how it’s dumped.
With different types in the enum
Now that the basics are covered, let’s use more complex enums. Let’s start to see what happens if we try to mix types in IntEnum and StrEnum. If Python can parse a type automatically in IntEnum it will do so:
class WithIntAndIntAsString(IntEnum): member1 = 1 # This will be parse automatically by Python. member2 = "10" print( "member1", type(WithIntAndIntAsString.member1.value), WithIntAndIntAsString.member1, ) print( "member2", type(WithIntAndIntAsString.member2.value), WithIntAndIntAsString.member2, )
member1 <class 'int'> 1 member2 <class 'int'> 10
Otherwise, it will raise a TypeError:
try: class MoreStringValues(StrEnum): member1 = "member1" # Python won’t cast anything to int. value2 = 10 except TypeError as e: print("Failing with error:", e)
Failing with error: 10 is not a string
If we need to mix types, we can use a base Enum:
class MultipleBasicTypes(Enum): member1 = "member1" member2 = 10 print( type(MultipleBasicTypes.member1.value), type(MultipleBasicTypes.member2.value), )
<class 'str'> <class 'int'>
If you need to enforce a consistent type within your enums, please refer to of my previous article and the enforce_member_type decorator.
Using enums with Pydantic models
Now that we know how enums behave with Pydantic, let’s dive the heart of this article: enums with Pydantic models as members and their usage in other Pydantic models! I’ll detail several possible strategies here and discuss the pros and cons of each one. You’ll probably have questions or remarks on these (and maybe even new methods), so don’t hesitate to use the comment section.
I’ll use an example to make things less abstract. I’ll leave your imagination work to find useful use-cases. You can also just see it as a way to have fun with Python and Pydantic.
Let’s imagine that you are running a blog platform. Each article on the platform is associated with a category. A category could be represented with this Pydantic model:
class Category(BaseModel): # The id of the category in the database id: int # A unique human-readable code identifying the category. code: str # The title to display to the users. title: str
Let’s say that for editorial reasons the list of categories is fixed. So this list can fit into an enum. The ids and the codes are set within the enum to ease referencing the categories in the codebase without needing to query the database. This will help to create new categories in the database when you add some. All you’ll have to do is loop over the members of the enum.
We’ll also assume it allows great simplification in you code to be able to hard code the id of the category directly. The enum makes this clean and readable. It will also make your JSON API more simple to use for your users. To create an article, they can use the human readable code of the category instead of the id. They can send this:
{"category": "programming", "title": "Fun with Pydantic"}
to get this back:
Article(category=ArticleCategory.programming, title="Fun with Pydantic")
No need to bother them with ids!
Let’s start this journey by creating the ArticleCategory enum:
class ArticleCategory(Enum): cooking = Category(id=1, code="cooking", title="Cooking") programming = Category(id=2, code="programming", title="Programming")
You’ll then be able to mention this in you article model:
class Article(BaseModel): # id of the article in the database. # Allow None in the model to enable user to create an article in the API. id: int | None = None category: ArticleCategory title: str
Note
Since ids don’t bring anything for this article, I’ll omit them from now on to make the examples a bit more simple.
How does this behave by default?
Let’s see:
article = Article.model_validate({ "category": ArticleCategory.programming, "title": "Fun with Pydantic", }) print(article)
id=None category=<ArticleCategory.programming: Category(id=2, code='programming', title='Programming')> title='Fun with Pydantic'
Our category field is a member of the enum as expected.
What about a call to model_dump?
print("Dump in Python mode:", article.model_dump(mode="python")) print("Dump in JSON mode:", article.model_dump(mode="json"))
Dump in Python mode: {'id': None, 'category': <ArticleCategory.programming: Category(id=2, code='programming', title='Programming')>, 'title': 'Fun with Pydantic'} Dump in JSON mode: {'id': None, 'category': {'id': 2, 'code': 'programming', 'title': 'Programming'}, 'title': 'Fun with Pydantic'}
If we dump for Python, the member of the enum is preserved and as JSON our model is serialized. So far, it’s only standard Pydantic stuff.
What if we pass a new Category instance? Let’s start with one equal to a member of the enum:
article = Article.model_validate({ "category": Category(id=2, code="programming", title="Programming"), "title": "Fun with Pydantic", }) print(article)
id=None category=<ArticleCategory.programming: Category(id=2, code='programming', title='Programming')> title='Fun with Pydantic'
Once again it works and category is a member of the enum.
What if we use a new category not part of the enum?
try: article = Article.model_validate({ "category": Category(id=3, code="travel", title="Travel"), "title": "Traveling", }) except ValidationError as e: print("It fails:") print(e)
It fails: 1 validation error for Article category Input should be Category(id=1, code='cooking', title='Cooking') or Category(id=2, code='programming', title='Programming') [type=enum, input_value=Category(id=3, code='travel', title='Travel'), input_type=Category] For further information visit https://errors.pydantic.dev/2.11/v/enum
As expected, we get an error because this category is not part of the enum.
And with a dict representing a category?
try: article = Article.model_validate({ "category": {"id": 2, "code": "programming", "title": "Programming"}, "title": "Fun with Pydantic", }) except ValidationError as e: print("It fails:") print(e)
It fails: 1 validation error for Article category Input should be Category(id=1, code='cooking', title='Cooking') or Category(id=2, code='programming', title='Programming') [type=enum, input_value={'id': 2, 'code': 'progra... 'title': 'Programming'}, input_type=dict] For further information visit https://errors.pydantic.dev/2.11/v/enum
That’s a bit sad. Likewise, we cannot reference our category by id or code when creating an article by default. Luckily, it’s a problem that can be solved. So, let’s do it!
Creating a custom Enum base class
If you read my previous article on custom classes and Pydantic, this idea should come to mind. It sure was the first solution that came to mine. All we need to to is create a BasePydanticEnum class inheriting from Enum and hook this enum to Pydantic with the __get_pydantic_core_schema__ class method. It will allow enum members to be loaded from their code field and serialized to their code field. I won’t detail __get_pydantic_core_schema__ too much here.
Note
To simplify code samples, I’ll assume all models used this way have a field named code for the serialization. I’ll provide at the end, as a bonus, a set of classes that allow the field to be specified for more flexibility. This way, you’ll be able to have multiple enums with Pydantic model as member and use different identifier fields.
class BasePydanticEnum(Enum): @classmethod def __get_pydantic_core_schema__( cls, _source: Any, _handler: GetCoreSchemaHandler ): # Get the member from the enum no matter what we have as input. # If we fail to find a matching member, it will fail. # It accepts: code, enum member and enum member value as input. get_member_schema = core_schema.no_info_plain_validator_function( cls._get_member ) return core_schema.json_or_python_schema( json_schema=get_member_schema, python_schema=get_member_schema, # Get the code from the value of the enum member. serialization=core_schema.plain_serializer_function_ser_schema( lambda member: member.value.code ), ) @classmethod def _get_member(cls, input_value: Any): for member in cls: # If the input is already a member or is a member value, let’s use it. if input_value == member or input_value == member.value: return member # If not, search for the member with input_value as code. if member.value.code == input_value: return member # Try to validate the input as a model, in case the user supplied a dict # representing a member. Validating during each loop is suboptimal, # improve this if you care about this feature. # Not easy since you can’t know easily the type of you member values by # default. Forcing the child to implement a _get_value_type class method # would solve this. try: model = type(member.value).model_validate(input_value) except ValidationError: continue else: # Validated successfully and matches the current member. if member.value == model: return member # Raise a ValueError if our search fails for Pydantic to create its proper # ValidationError. raise ValueError(f"Failed to convert {input_value} to a member of {cls}") class BasePydanticEnumArticleCategory(BasePydanticEnum): cooking = Category(id=1, code="cooking", title="Cooking") programming = Category(id=2, code="programming", title="Programming") class BasePydanticEnumArticle(BaseModel): category: BasePydanticEnumArticleCategory title: str
Let’s check it works as expected. I’ll just provide the full code examples with their output, since it should be straightforward at this point.
article = BasePydanticEnumArticle.model_validate( {"category": BasePydanticEnumArticleCategory.programming, "title": "Fun with Pydantic"} ) print("Creation from member:", article) print("Dump in Python mode:", article.model_dump(mode="python")) print("Dump in JSON mode:", article.model_dump(mode="json")) print( "From a category object that’s equal to a member:", BasePydanticEnumArticle.model_validate( { "category": Category( id=2, code="programming", title="Programming" ), "title": "Fun with Pydantic", } ), ) print( "Creating from a valid code:", BasePydanticEnumArticle.model_validate( {"category": "programming", "title": "Fun with Pydantic"} ), ) print( "Creating from a valid dict:", BasePydanticEnumArticle.model_validate( { "category": {"id": 2, "code": "programming", "title": "Programming"}, "title": "Fun with Pydantic", } ), ) try: article = BasePydanticEnumArticle.model_validate( { "category": Category( id=3, code="travel", title="Travel" ), "title": "Traveling", } ) except ValidationError: print("Creating with a non member category fails as expected.") try: article = BasePydanticEnumArticle.model_validate( {"category": "travel", "title": "Traveling"} ) except ValidationError: print("Creating from an invalid code fails as expected.")
Creation from member: category=<BasePydanticEnumArticleCategory.programming: Category(id=2, code='programming', title='Programming')> title='Fun with Pydantic' Dump in Python mode: {'category': 'programming', 'title': 'Fun with Pydantic'} Dump in JSON mode: {'category': 'programming', 'title': 'Fun with Pydantic'} From a category object that’s equal to a member: category=<BasePydanticEnumArticleCategory.programming: Category(id=2, code='programming', title='Programming')> title='Fun with Pydantic' Creating from a valid code: category=<BasePydanticEnumArticleCategory.programming: Category(id=2, code='programming', title='Programming')> title='Fun with Pydantic' Creating from a valid dict: category=<BasePydanticEnumArticleCategory.programming: Category(id=2, code='programming', title='Programming')> title='Fun with Pydantic' Creating with a non member category fails as expected. Creating from an invalid code fails as expected.
So far so good. And problem solved for all the use cases! But, this solution also has a set of drawbacks:
- All your enums must inherit from a custom enum class.
- All the logic is in this custom class while the logic for the serialization and deserialization is about the model.
- It requires advanced concepts.
- It’s complex. And to make it generic it’s even more complex since part of the logic must be spread: the serialization/deserialization logic stays in the base enum class and the name of the identifier field must go into the model or the child class. It’s easier to enforce in the child enum since we can create the method in the base class with a raise NotImplemented body.
- Advanced note: since enum already has a metaclass, the custom class cannot inherit from ABCMeta and use the @abstract_method decorator.
- This serialization/deserialization is forced on all users.
Can we find a better way to do this? Please read-on!
Alternative with annotations
If you’ve used Pydantic, you may wander if we can use the Annotate pattern to do this. Using BeforeValidator and PlainSerializer we can. It’s verbose and error prone (we must specify both each time we define a field). It also relies on less advanced Pydantic concepts.
The serializer is easy and reusable:
EnumMemberToCodeSerializer = PlainSerializer(lambda member: member.value.code)
The validator a bit less:
def get_member(enum_cls: type[Enum], input_value: Any): # Same logic that in the class. I left the loading a dict into a model # out for simplicity. for member in enum_cls: if input_value == member or input_value == member.value: return member if member.value.code == input_value: return member raise ValueError(f"Failed to find a member in {enum_cls} from {input_value}") # The actual validator must be built with the enum as input so it can access # the proper members. def create_enum_from_code_validator(enum_cls: type[Enum]): return BeforeValidator(lambda input_value: get_member(enum_cls, input_value))
Now let’s glue it together in a model and test how this behaves:
class AnnotatedArticle(BaseModel): category: Annotated[ ArticleCategory, create_enum_from_code_validator(ArticleCategory), EnumMemberToCodeSerializer, ] m = AnnotatedArticle.model_validate({"category": "programming"}) print("Checking behavior:") print(m) print(m.model_dump(mode="python")) print(m.model_dump(mode="json")) print(AnnotatedArticle.model_validate({"category": ArticleCategory.cooking})) print(AnnotatedArticle.model_validate({ "category": ArticleCategory.programming.value, })) try: AnnotatedArticle.model_validate({"category": "travel"}) assert False, "Should not go there" except ValidationError: print("Failed with invalid code")
Checking behavior: category=<ArticleCategory.programming: Category(id=2, code='programming', title='Programming')> {'category': 'programming'} {'category': 'programming'} category=<ArticleCategory.cooking: Category(id=1, code='cooking', title='Cooking')> category=<ArticleCategory.programming: Category(id=2, code='programming', title='Programming')> Failed with invalid code
Note
You may wander why I didn’t use something like category: build_pydantic_enum_type_annotation(ArticleCategoryForAnnotated) to build everything with one function. build_pydantic_enum_type_annotation could be defined like so:
def build_pydantic_enum_type_annotation(enum_cls: type[Enum]): return Annotated[ enum_cls, BeforeValidator(lambda input_value: get_member(enum_cls, input_value)), EnumMemberToCodeSerializer, ]
The reason is: it won’t work. It requires to use a call expression with a type expression and that’s not allowed.
It works, but is also quite weird and complex. I think I prefer the first solution. The field must be annotated with both a custom serializer and a custom validator or it won’t work. At least, it only concerns the leaf model, the rest can remain standard. I’m still sure we can do better.
Logic into the model that’s the member of the enum
Let’s try another approach in which the enum stays standard and it’s its Pydantic model members which know how to serialize themselves. In this solution:
- Serialization is made easy with a model_serializer placed on the enum that is a member of the class.
- Derialization is harder:
- We cannot hook into the model because it won’t be used: we will receive a string instead of an enum value and this string won’t be equal to any member of the enum. So Pydantic will stop there and won’t use any model_validator(mode="before") we may have to load the data. We don’t have an easy access to the enum there anyway to identify the member to use.
- We cannot use a model_validator on the enum since it’s not a Pydantic model. It would also mean we have part of the logic in the model and part on the enum which breaks locality. It would go against what we are trying to achieve here.
- We could put a model_validator on the leaf model but we break locality even further.
- What works is to hook into the __eq__ method of the model: if the value we receive is the identifier of the model, we can let Pydantic know it found the right member of the enum. When it loads the data, Pydantic will loop over the enum member and check to see whether its input is equal to one of their value. If so, it will use it.
class BaseComplexEnumModel(BaseModel): @model_serializer() def _serialize_model(self): return self.code def __eq__(self, other) -> bool: # Both are Pydantic models, let's see whether Pydantic thinks they are # equal by using the BaseModel.__eq__ if isinstance(other, BaseModel): return super().__eq__(other) # We have an identifier, let's see whether it matches this model's one. return self.code == other class ComplexCategory(BaseComplexEnumModel): code: str title: str class ComplexArticleCategory(Enum): programming = ComplexCategory(code="programming", title="Programming") cooking = ComplexCategory(code="cooking", title="Cooking")
Let’s run some basic tests:
class ComplexArticle(BaseModel): category: ComplexArticleCategory m = ComplexArticle.model_validate({"category": "programming"}) print("Checking behavior:", m) print("From a member value:", ComplexArticle.model_validate({ "category": ComplexArticleCategory.programming.value, })) print("From a member:", ComplexArticle.model_validate({ "category": ComplexArticleCategory.cooking, })) try: ComplexArticle.model_validate({"category": "travel"}) assert False, "Should not go there" except ValidationError: print("Failed with invalid code") # weird but it works since we overloaded __eq__ to load that. assert ComplexCategory(code="some_code", title="Some title") == "some_code"
Checking behavior: category=<ComplexArticleCategory.programming: ComplexCategory(code='programming', title='Programming')> From a member value: category=<ComplexArticleCategory.programming: ComplexCategory(code='programming', title='Programming')> From a member: category=<ComplexArticleCategory.cooking: ComplexCategory(code='cooking', title='Cooking')> Failed with invalid code
I find this solution very weird and magical. And it leaks outside the Pydantic context by changing deeply how equality behaves. It works, it illustrates a “creative” use of overloading __eq__ for advanced needs, but I can’t remand it. Once again, I’m sure we can find something better.
Putting the logic in the leaf model
After trying to put all the logic in the model that’s used in the enum, let’s try to put it at the other end: in the model that uses the enum. Just like with the annotated example, this solution uses a custom serializer and a custom validator.
class BaseLeafModelWithCodeSupport(BaseModel): @model_validator(mode="before") @classmethod def _load_pydantic_enum(cls, values): for field_name, field_def in cls.model_fields.items(): # Do nothing for fields whose are not defined as enums if field_name not in values or not issubclass(field_def.annotation, Enum): continue # Loop over the enum members now that we know that field_def.annotation # is an enum. for member in field_def.annotation: # Find the allowed values as input: the member itself, the member # value and the code. # If one of them matches, member is the member we are looking for. if values[field_name] in cls._get_allowed_input_values_for_member( member ): values[field_name] = member return values @classmethod def _get_allowed_input_values_for_member(cls, member): # It’s an enum member without a code. Can’t do anything with it. # Let Pydantic load it its usual way. if not hasattr(member.value, "code"): return (member, member.value) return (member, member.value, member.value.code) @model_serializer(when_used="json", mode="wrap") def _serialize_pydantic_enums( self, handler: SerializerFunctionWrapHandler, info: SerializationInfo, ): # Let Pydantic serialize the values automatically. values = handler(self) for field_name, field_def in self.__pydantic_fields__.items(): # Find the fields defined as enums. if not issubclass(field_def.annotation, Enum): continue # A field may not have a serialization alias, use its name in # this case instead. serialization_alias = field_def.serialization_alias or field_name # Use the serialization alias, only when dumping by alias. serialization_field_name = ( serialization_alias if info.by_alias else field_name ) field_value = getattr(self, field_name) # Not an enum member with a code. We don’t care about it. if not hasattr(field_value, "code"): continue # Some fields may be excluded from serialization (because they are # unset for instance). Let's ignore those. if serialization_field_name in values: values[serialization_field_name] = field_value.code return values
Let’s test:
class ArticleWithCodeSupport(BaseLeafModelWithCodeSupport): category: ArticleCategory category_with_alias: ArticleCategory = Field(serialization_alias="cat", default=ArticleCategory.programming) m = ArticleWithCodeSupport.model_validate({"category": "programming", "category_with_alias": "programming"}) print("Checking behaviors") print(m) print(m.model_dump(mode="python")) print(m.model_dump(mode="json", by_alias=True)) print(m.model_dump(mode="json", by_alias=False)) print(ArticleWithCodeSupport.model_validate({"category": ArticleCategory.cooking})) print(ArticleWithCodeSupport.model_validate({"category": ArticleCategory.programming.value})) try: ArticleWithCodeSupport.model_validate({"category": "travel"}) assert False, "Should not go there" except ValidationError: print("Failed with invalid code")
Checking behaviors category=<ArticleCategory.programming: Category(id=2, code='programming', title='Programming')> category_with_alias=<ArticleCategory.programming: Category(id=2, code='programming', title='Programming')> {'category': <ArticleCategory.programming: Category(id=2, code='programming', title='Programming')>, 'category_with_alias': <ArticleCategory.programming: Category(id=2, code='programming', title='Programming')>} {'category': {'id': 2, 'code': 'programming', 'title': 'Programming'}, 'cat': {'id': 2, 'code': 'programming', 'title': 'Programming'}} {'category': {'id': 2, 'code': 'programming', 'title': 'Programming'}, 'category_with_alias': {'id': 2, 'code': 'programming', 'title': 'Programming'}} category=<ArticleCategory.cooking: Category(id=1, code='cooking', title='Cooking')> category_with_alias=<ArticleCategory.programming: Category(id=2, code='programming', title='Programming')> category=<ArticleCategory.programming: Category(id=2, code='programming', title='Programming')> category_with_alias=<ArticleCategory.programming: Category(id=2, code='programming', title='Programming')> Failed with invalid code
Works very well! It’s based on model_validator and model_serializer which are more common than __get_pydantic_core_schema__. It requires care in the implementation to work correctly and to prevent issues with “normal” enum fields. Most notably, it must loop over the field definitions and access advanced attributes of the model.
On the bright side, all the logic is in the model that actually cares about it (well, in its base class to be exact). I think it’s the best solution since everything is done locally in one place and it only impacts the model that requires the behavior regarding the code attribute. The implementation is not too complex if properly documented (at least in my opinion).
Conclusion
Which solution would I choose? The “logic in the leaf model with model serialize and model validator” since I think it’s the cleanest and one of the more simple. It’s also the one with the least impact on other part of the code: only the model that need the weird “code for member behavior” is impacted. I may also consider the first solution that relies on __get_pydantic_core_schema__ since I think it’s quite clean even if logic is in a weird place. I’d probably do this if I have many custom classes and thus, __get_pydantic_core_schema__ has become very familiar to me.
I’d just discard the other possibilities as too weird or complex.
I hope you enjoyed this article about a weird use case of Pydantic with enums as much as I did and learned things on the way. I myself really liked the exploration and the different ideas that required each case to work. Don’t hesitate to leave a comment below if needed!
Bonus
In all the examples above, I relied on a field named code to do the serialization/deserialization. I didn’t try to support any other field. Here, I’ll give ways for the 1st and 4th method to select which field must be used instead. I’ll assume you’ve read the article and are at ease with it so I won’t comment the code as much. If you have a problem, let me know in the comments!
Note
I haven’t tested these solutions as much so they may be a bit buggy.
class EnumModel(BaseModel): code: str class BasePydanticEnumWithIdentifier(Enum): """Allow enum member to be loaded from their identifier (a field of the model) and then serialized to their identifier. """ @classmethod def __get_pydantic_core_schema__(cls, _source: Any, _handler: GetCoreSchemaHandler): # Get the member from the enum no matter what we have as input. If we fail to find a matching member, it will fail. # It accepts: code, enum member and enum member value as input. get_member_schema = core_schema.no_info_plain_validator_function(cls._get_member) return core_schema.json_or_python_schema( json_schema=get_member_schema, python_schema=get_member_schema, serialization=core_schema.plain_serializer_function_ser_schema(cls._serialize_to_identifier), ) @classmethod def _get_member(cls, input_value: Any): for member in cls: if input_value == member or input_value == member.value: return member if getattr(member.value, cls._get_identifier()) == input_value: return member raise ValueError(f"Failed to convert {input_value} to a member of {cls}") @classmethod def _serialize_to_identifier(cls, member) -> Any: return getattr(member.value, cls._get_identifier()) @classmethod def _get_identifier(cls) -> str: # Used to override in each enum what field of the member to use to load/serialize the data. # Note: using @abstractmethod is NOT straightforward: both Enum and ABC uses metaclass and they are not compatible. # We should be able to create our own to adapt it. Probably not worth it. raise NotImplementedError class ModelValues(BasePydanticEnum): first_model = EnumModel(code="first_model") second_model = EnumModel(code="second_model") @classmethod def _get_identifier(cls): return "code" class MyModel(BaseModel): model_value: ModelValues m = MyModel.model_validate({"model_value": "first_model"}) print(m) print(m.model_dump())
model_value=<ModelValues.first_model: EnumModel(code='first_model')> {'model_value': 'first_model'}
class EnumModel(BaseModel): code: str @property def identifier(self): # Sad that it must be on this model and on no the leaf. return "code" class ModelDecoratorModelValues(Enum): first_choice = EnumModel(code="first_choice") second_choice = EnumModel(code="second_choice") class BaseModelDecoratorWithComplexEnum(BaseModel): @model_validator(mode="before") @classmethod def _load_pydantic_enum(cls, values): for field_name, field_def in cls.model_fields.items(): if field_name not in values or not issubclass(field_def.annotation, Enum): continue # Loop over the enum members now that we know that field_def.annotation is an enum. for member in field_def.annotation: if values[field_name] in cls._get_allowed_input_values_for_member(member): values[field_name] = member return values @classmethod def _get_allowed_input_values_for_member(cls, member): if not hasattr(member.value, "identifier"): return (member, member.value) member_identifier_field_name = member.value.identifier member_identifier = getattr(member.value, member_identifier_field_name) return (member, member.value, member_identifier) @model_serializer(when_used="json", mode="wrap") def _serialize_pydantic_enums(self, handler: SerializerFunctionWrapHandler, info: SerializationInfo): values = handler(self) for field_name, field_def in self.model_fields.items(): if not issubclass(field_def.annotation, Enum): continue # A field may not have a serialization alias, use its name in this case instead. serialization_alias = field_def.serialization_alias or field_name serialization_field_name = serialization_alias if info.by_alias else field_name # Some fields may be excluded from serialization (because they are unset for instance). Let's ignore those. if serialization_field_name in values: values[serialization_field_name] = self._get_allowed_input_values_for_member(getattr(self, field_name)) return values class ModelDecoratorWithComplexEnum(BaseModelDecoratorWithComplexEnum): model_value: ModelValues enum_value: ModelDecoratorModelValues m = MyModel.model_validate({"model_value": "first_model"}) print(m) print(m.model_dump())
model_value=<ModelValues.first_model: EnumModel(code='first_model')> {'model_value': 'first_model'}