Python 枚举和 Pydantic ：接受枚举成员的组合

Question

我有一个enum:

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">

我不能让这个字段作为一个带有验证器的字符串字段，因为我使用了一个测试库(假设和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

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

如何将此字段保持为Enum类型(使我的模拟结构仍然有效)并使pydantic同时接受复合值?

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

Answer 1

我进一步研究了这个问题，我相信这样的东西可能会有所帮助。您可以创建一个新类来定义该属性，该属性是枚举值的列表。

这个类可以提供一个自定义的validate方法，并提供一个__modify_schema__来保持json模式中字符串的信息。

可以像这样为连接枚举的泛型列表定义一个基类:

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"
]
}