为列表中的不同 JSON 模式模型定义约束

我有一些 JSON，其结构类似于下图所示。 threshold 列表表示类型可以是 "type": "upper_limit" 或 "type": "range" 的对象。请注意， "target" 值应该是整数或浮点数，具体取决于对象的类型。

{
    "name": "blah",
    "project": "blah blah",
    "threshold": [
        {
            "id": "234asdflkj",
            "group": "walkers",
            "type": "upper_limit",
            "target": 20,
            "var": "distance"
        },
        {
            "id": "asdf34asf2654",
            "group": "runners",
            "type": "range",
            "target": 1.7,
            "var": "speed"
        }
    ]
}

为上述数据生成 JSON 模式的 Pydantic 模型如下所示：

class ThresholdType(str, Enum):
    upper_limit = "upper_limit"
    range = "range"


class ThresholdUpperLimit(BaseModel):
    id: str
    group: str
    type: ThresholdType = "upper_limit"
    target: int = Field(gt=2, le=20)
    var: str


class ThresholdRange(BaseModel):
    id: str
    group: str
    type: ThresholdType = "range"
    target: float = Field(gt=0, lt=10)
    var: str


class Campaign(BaseModel):
    name: str
    project: str
    threshold: list[ThresholdUpperLimit | ThresholdRange]

模型验证 JSON，但有约束对于该类型， target 值将被忽略。例如，如果阈值对象包含 "type": "range", "target": 12, ，则不会抛出错误，因为它被解析为整数，因此使用 ThresholdUpperLimit 定义的约束；但应使用 ThresholdRange 定义的约束。关于如何正确处理这个问题有什么建议吗？

是对的，的 Pydantic 模型无法按预期捕获 target 字段的类型约束，因为它取决于 type 字段。这是因为在使用 Union 时，Pydantic 首先会尝试使用第一个匹配的模型（在本例中为 ThresholdUpperLimit ），而没有考虑类型约束。

可以通过以下几种方式解决此问题：

1. 使用自定义验证器和根验证器：

• 从 Union 中删除 type 字段，并在根验证器中进行处理。
• 使用 root_validator 根据 type 字段的值对 target 字段应用验证逻辑。

from pydantic import BaseModel, Field, root_validator, ValidationError


class ThresholdType(str, Enum):
    upper_limit = "upper_limit"
    range = "range"


class ThresholdBase(BaseModel):
    id: str
    group: str
    type: ThresholdType
    target: Union[int, float]
    var: str

    @root_validator
    def validate_target_type(cls, values):
        target_type = values.get("type")
        target_value = values.get("target")
        if target_type == ThresholdType.upper_limit:
            if not isinstance(target_value, int) or not 2 < target_value <= 20:
                raise ValueError("target must be an integer between 2 and 20 for upper_limit type")
        elif target_type == ThresholdType.range:
            if not isinstance(target_value, float) or not 0 < target_value < 10:
                raise ValueError("target must be a float between 0 and 10 for range type")
        return values


class Campaign(BaseModel):
    name: str
    project: str
    threshold: list[ThresholdBase]

2. 使用 Discriminated Unions（Pydantic >= 1.0）：

从 Pydantic 1.0 开始，可以使用 discriminator 字段创建区分联合。这允许 Pydantic 根据 type 字段的值选择正确的模型。

from pydantic import BaseModel, Field, discriminator


class ThresholdType(str, Enum):
    upper_limit = "upper_limit"
    range = "range"


class ThresholdBase(BaseModel):
    id: str
    group: str
    type: ThresholdType
    target: Union[int, float]
    var: str


class ThresholdUpperLimit(ThresholdBase):
    type: ThresholdType = "upper_limit"
    target: int = Field(gt=2, le=20)


class ThresholdRange(ThresholdBase):
    type: ThresholdType = "range"
    target: float = Field(gt=0, lt=10)


class Campaign(BaseModel):
    name: str
    project: str
    threshold: list[ThresholdBase] = Field(..., discriminator='type')

这两种方法都能正确验证 target 字段，具体取决于 type 字段的值。选择哪种方法取决于的偏好和 Pydantic 版本。