validation - Pydantic 与枚举子类？

Question

想象一下我有类似的东西

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

然后就有了

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

然后，我有一个 pydantic 模型:

class MyModel(BaseModel):
    my_field: SpecialEnum

我希望 pydantic 验证 my_field 是 SpecialEnum 子类的某个实例。 然而，pydantic 使用其特殊的枚举验证器，并使用 SpecialEnum(field_value) 进行验证，从而导致始终出现 ValidationError - 因为 SpecialEnum 的目的只是提供功能或界面，而不是值。

有什么解决办法吗？

Answer 1

无需解决方法。您可以通过将特殊的 __get_validators__ 类方法添加到基类 SpecialEnum 来覆盖默认验证。这是一个完整的工作示例:

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`

不过这里有一个警告。通常，您可以定义一些附加逻辑来接受枚举成员的值，甚至允许成员的名称，并让验证器返回相应的枚举成员本身。 (最后一行中的 wrong2 示例。)

但是，由于您的验证方法必须在基类 SpecialEnum 上定义，因为这是您用来注释 my_field 的方法，因此无法知道用于检查值/名称的枚举子类。

如果您对此感兴趣，您也许可以为此构造一些额外的逻辑，但我认为只有在确保 SpecialEnum 的所有子类的所有枚举成员是不同的。

希望这有帮助。