在 pydantic 模型中包含可能的枚举值

Question

我有一个包含枚举值的模型 ModelWithEnum。除了该值之外，我还希望模型从该枚举中输出所有可能的值（这些枚举类似于范围，例如 A、B、C；A、B、C、D 等）。

我制作了 EnumModel 一个 GenericModel，因为我不想为每个 Enum class 创建一个新模型。我正在使用仿制药迈出第一步，但不确定我是否在正确的轨道上。

这是我想要实现的目标

# `ModelWithEnum` holds a `enum` field of type `EnumsAC`,
# so possible enum values would be "A", "B", "C"
model = ModelWithEnum(enum={"value": "A"})

assert model.enum.value == "A"
assert model.enum.possible_values == ["A", "B", "C"]

完整示例代码：

from enum import Enum
from typing import Generic, List, Optional, TypeVar
from pydantic.generics import GenericModel
from pydantic import BaseModel

EnumT = TypeVar("EnumT")


class EnumsAC(str, Enum):
    A = "A"
    B = "B"
    C = "C"


class EnumsAE(str, Enum):
    """Another Enums class"""

    A = "A"
    B = "B"
    C = "C"
    D = "D"
    E = "E"


class EnumModel(GenericModel, Generic[EnumT]):
    value: EnumT
    # How can I auto-fill this list with all possible values from EnumT
    # like if I hardcoded e.g. `list(EnumsAC)`
    possible_values: List[str] = []


class ModelWithEnum(BaseModel):
    enum: EnumModel[EnumsAC]

# `ModelWithEnum` holds a `enum` field of type `EnumsAC`,
# so possible enum values would be "A", "B", "C"
model = ModelWithEnum(enum={"value": "A"})

assert model.enum.value == "A" # ok
assert model.enum.possible_values == ["A", "B", "C"] # fails

Answer 1

您可以使用 root_validator 来自动完成可能的值。以您的模型为例：

class EnumModel(GenericModel, Generic[EnumT]):
    value: EnumT
    possible_values: List[str] = []

    class Config:
        validate_assignment = True

    @root_validator
    def root_validate(cls, values):
        values["possible_values"] = [item for item in values['value'].__class__]
        return values

在 pydantic 模型中包含可能的枚举值

Include possible enum values in pydantic model

python

generics

enums

pydantic

这是我想要实现的目标

完整示例代码：