Pydantic with enum subclasses?-CodePudding

So imagine I have something like

class SpecialEnum(Enum):
    def do_stuff(self):
         # some cool operation

And then have

class MySpecialEnum(SpecialDescriptorEnum):
   A = 1
   B = 2

Then, I have a pydantic model with:

class MyModel(BaseModel):
    my_field: SpecialEnum

I want pydantic to validate that my_field is some instance of a SpecialEnum subclass. However, pydantic uses its special enum validator instead, and validates with SpecialEnum(field_value) instead, resulting in what will always be a ValidationError - since SpecialEnum was only ever meant to provide functionality or interface, and not values.

Is there any workaround for that?

CodePudding user response：

No need for a workaround. You can override the default validation by adding the special __get_validators__ class method to your base SpecialEnum class. Here is a full working example:

from __future__ import annotations
from collections.abc import Callable, Iterator
from enum import Enum
from typing import Any

from pydantic import BaseModel


class SpecialEnum(Enum):
    @classmethod
    def __get_validators__(cls) -> Iterator[Callable[..., Any]]:
        yield cls.validate

    @classmethod
    def validate(cls, v: SpecialEnum) -> SpecialEnum:
        if not isinstance(v, cls):
            raise TypeError(f"Must be an instance of `{cls.__name__}`")
        return v

    def do_stuff(self) -> None:
        print("stuff", self)


class MySpecialEnum(SpecialEnum):
    A = 1
    B = 2


class MyModel(BaseModel):
    my_field: SpecialEnum


class Unrelated(Enum):
    foo = "foo"


if __name__ == '__main__':
    instance = MyModel(my_field=MySpecialEnum.A)
    instance.my_field.do_stuff()
    # wrong1 = MyModel(my_field=Unrelated.foo)  # causes a `ValidationError`
    # wrong2 = MyModel(my_field=1)  # ALSO causes a `ValidationError`

There is a caveat here though. Usually you would be able to define some additional logic to accept the enum member's value or even allow the member's name and have the validator return the corresponding enum member itself. (The wrong2 example in the last line.)

But since your validation methods must be defined on the base class SpecialEnum because that is the one you use to annotate my_field, there is no way to know which enum subclass to check the value/name against.

If that is of interest to you, you might be able to construct some additional logic for that, but I think this will only be possible if you ensure that all enum members of all subclasses of SpecialEnum are distinct.

Hope this helps.