Python Enum 和 Pydantic : 接受 enum 成员的组合

Question

我有一个枚举：

from enum import Enum

class MyEnum(Enum):
    val1 = "val1"
    val2 = "val2"
    val3 = "val3"

我想根据该枚举验证一个 pydantic 字段。

from pydantic import BaseModel

class MyModel(BaseModel):
    my_enum_field: MyEnum

但我希望此验证也接受由 Enum 成员组成的字符串。

例如：“val1_val2_val3”或“val1_val3”是有效输入。

我无法将此字段作为带有验证器的字符串字段，因为我使用需要此类型的测试库 (hypothesis and pydantic-factories) 以呈现枚举中的值之一（用于模拟随机输入)

所以这个 :

from pydantic import BaseModel, validator

class MyModel(BaseModel):
    my_enum_field: str

    @validator('my_enum_field', pre=True)
    def validate_my_enum_field(cls, value):
        split_val = str(value).split('_')
        if not all(v in MyEnum._value2member_map_ for v in split_val):
            raise ValueError()
        return value

可以工作，但破坏了我的测试套件，因为该字段不再是枚举类型。

如何将此字段保留为枚举类型（以使我的模拟结构仍然有效）并使 pydantic 同时接受复合值？

到目前为止，我尝试动态扩展枚举，但没有成功。

Answer 1

我对此进行了更深入的研究，我相信类似的内容可能会有所帮助。您可以创建一个新的 class 来定义作为枚举值列表的属性。

此 class 可以提供自定义的 validate 方法并提供 __modify_schema__ 以保持有关在 json 架构中作为字符串的信息。

我们可以为这样的串联枚举的通用列表定义一个基础 class：

from typing import Generic, TypeVar, Type
from enum import Enum

T = TypeVar("T", bound=Enum)


class ConcatenatedEnum(Generic[T], list[T]):
    enum_type: Type[T]

    @classmethod
    def __get_validators__(cls):
        yield cls.validate

    @classmethod
    def validate(cls, value: str):
        return list(map(cls.enum_type, value.split("_")))

    @classmethod
    def __modify_schema__(cls, field_schema: dict):
        all_values = ', '.join(f"'{ex.value}'" for ex in cls.enum_type)
        field_schema.update(
            title=f"Concatenation of {cls.enum_type.__name__} values",
            description=f"Underscore delimited list of values {all_values}",
            type="string",
        )
        if "items" in field_schema:
            del field_schema["items"]

在 __modify_schema__ 方法中，我还提供了一种方法来生成对哪些值有效的描述。

要在您的应用程序中使用它：

class MyEnum(Enum):
    val1 = "val1"
    val2 = "val2"
    val3 = "val3"


class MyEnumList(ConcatenatedEnum[MyEnum]):
    enum_type = MyEnum


class MyModel(BaseModel):
    my_enum_field: MyEnumList

示例模型：

print(MyModel.parse_obj({"my_enum_field": "val1"}))
print(MyModel.parse_obj({"my_enum_field": "val1_val2"}))

my_enum_field=[<MyEnum.val1: 'val1'>]
my_enum_field=[<MyEnum.val1: 'val1'>, <MyEnum.val2: 'val2'>]

示例架构：

print(json.dumps(MyModel.schema(), indent=2))

{
  "title": "MyModel",
  "type": "object",
  "properties": {
    "my_enum_field": {
      "title": "Concatenation of MyEnum values",
      "description": "Underscore delimited list of values 'val1', 'val2', 'val3'",
      "type": "string"
    }
  },
  "required": [
    "my_enum_field"
  ]
}

Python Enum 和 Pydantic : 接受 enum 成员的组合

Python Enum and Pydantic : accept enum member's composition

python

enums

metaprogramming

pydantic