如何在Pydantic中实现依building_type的条件必填字段校验？

问题说明 #

需要校验两组JSON数据，要求：

• 两组JSON均能通过Pydantic模型校验
• 仅当building_type取值为flat时，apartment字段为必填项

待校验JSON1：

{ 'building_type' : 'flat',
  'direction': 'Fake Street',
  'number': 123,
  'apartment': 3
}

待校验JSON2：

{ 'building_type' : 'house',
  'direction': 'Fake Street',
  'number': 123
}

现有基础Pydantic模型：

from enum import Enum
from pydantic import BaseModel


class Building(Enum):
  FLAT = 'flat'
  HOUSE = 'house'

class Address(BaseModel):
  model_config = ConfigDict(extra='forbid')
  building_type: Building
  direction: str

实现方案 #

可以通过Pydantic的**模型校验器（model_validator）**结合可选字段来实现需求，具体修改后的模型代码如下：

from enum import Enum
from pydantic import BaseModel, Field, model_validator
from typing import Optional


class Building(Enum):
    FLAT = 'flat'
    HOUSE = 'house'

class Address(BaseModel):
    model_config = ConfigDict(extra='forbid')
    building_type: Building
    direction: str
    number: int  # 补充原JSON中存在的number字段
    apartment: Optional[int] = Field(None)  # 将apartment设为可选，默认值为None

    @model_validator(mode='after')
    def check_apartment_required(self) -> 'Address':
        # 当building_type为FLAT时，校验apartment是否存在
        if self.building_type == Building.FLAT and self.apartment is None:
            raise ValueError("building_type为'flat'时，apartment字段为必填项")
        return self

代码说明 #

1. 补充字段：原模型缺少number字段，需添加以匹配待校验JSON的结构，否则会因extra='forbid'触发校验错误。
2. 可选字段设置：将apartment定义为Optional[int]并设置默认值None，确保house类型的地址可以不传入该字段。
3. 条件校验逻辑：使用model_validator(mode='after')在模型实例化完成后进行校验，判断当building_type为FLAT时，apartment是否为有效值，若为空则抛出错误。

测试验证 #

• 验证合法的flat类型地址：

address_flat = Address(**{
    'building_type': 'flat',
    'direction': 'Fake Street',
    'number': 123,
    'apartment': 3
})
# 校验通过，无报错

• 验证合法的house类型地址：

address_house = Address(**{
    'building_type': 'house',
    'direction': 'Fake Street',
    'number': 123
})
# 校验通过，无报错

• 验证非法的flat类型地址（缺少apartment）：

try:
    Address(**{
        'building_type': 'flat',
        'direction': 'Fake Street',
        'number': 123
    })
except ValueError as e:
    print(e)  # 输出：building_type为'flat'时，apartment字段为必填项

内容的提问来源于stack exchange，提问作者vll1990