如何让Pydantic根据字典结构自动实例化对应子类BaseModel

需要将不同结构的Python字典转换为对应的Pydantic BaseModel实例。现有代码定义了ConditionType枚举、BaseConditionalProps的两个子类ConditionalExpressionProps、CycleDurationTrendProps，以及包含这两个子类的ConditionalConfig和SendNotificationChannel模型。期望通过SendNotificationChannel(**dict)实例化时，根据字典中的cond_type自动将conditional_config的condition_prop实例化为对应子类（如ConditionalExpressionProps或CycleDurationTrendProps），但当前实例化后condition_prop始终是父类BaseConditionalProps，需要解决这个问题。

现有代码 #

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

class ConditionType(str, Enum):
    EXPRESSION = 'EXPRESSION'
    CYCLE_DUR_TREND = 'CYCLE_DUR_TREND'

class BaseConditionalProps(BaseModel):
    pass  # 父类为基础空模型

class ConditionalExpressionProps(BaseConditionalProps):
    conditional_expression: str

class CycleDurationTrendProps(BaseConditionalProps):
    direction_up : bool = True
    n : int = Field(1, ge=1, le=1000)

class ConditionalConfig(BaseModel):
    cond_type: ConditionType = ConditionType.EXPRESSION
    condition_prop: BaseConditionalProps

class SendNotificationChannel(BaseModel):
    id: str
    customer_id: str
    conditional_config: Optional[ConditionalConfig]

测试示例 #

>>> channel = {"id": "1", "customer_id": "c1", "conditional_config" : {"cond_type": "EXPRESSION", "condition_prop":{"conditional_expression":"ALPHA"}}}
>>> channel_obj = SendNotificationChannel(**channel)
>>> channel_obj
SendNotificationChannel(id='1', customer_id='c1', conditional_config=ConditionalConfig(cond_type=<ConditionType.EXPRESSION: 'EXPRESSION'>, condition_prop=BaseConditionalProps()))
>>> channel_obj.conditional_config.condition_prop
BaseConditionalProps()

方案1：Pydantic v2 模型验证器（精准匹配） #

利用Pydantic v2的model_validator，根据cond_type字段手动将condition_prop映射到对应子类：

from pydantic import model_validator

class ConditionalConfig(BaseModel):
    cond_type: ConditionType = ConditionType.EXPRESSION
    condition_prop: ConditionalExpressionProps | CycleDurationTrendProps

    @model_validator(mode='before')
    def match_condition_prop(cls, values):
        cond_type = values.get('cond_type')
        prop_data = values.get('condition_prop')
        if cond_type == ConditionType.EXPRESSION:
            values['condition_prop'] = ConditionalExpressionProps(**prop_data)
        elif cond_type == ConditionType.CYCLE_DUR_TREND:
            values['condition_prop'] = CycleDurationTrendProps(**prop_data)
        return values

测试验证：

>>> channel_obj = SendNotificationChannel(**channel)
>>> channel_obj.conditional_config.condition_prop
ConditionalExpressionProps(conditional_expression='ALPHA')

方案2：Pydantic v2 鉴别联合类型（规范写法） #

通过给父类添加鉴别字段，让Pydantic自动匹配子类，需要调整输入结构将类型标识合并到condition_prop中：

from pydantic import Discriminator, Field

class BaseConditionalProps(BaseModel):
    type: ConditionType

class ConditionalExpressionProps(BaseConditionalProps):
    type: ConditionType = ConditionType.EXPRESSION
    conditional_expression: str

class CycleDurationTrendProps(BaseConditionalProps):
    type: ConditionType = ConditionType.CYCLE_DUR_TREND
    direction_up : bool = True
    n : int = Field(1, ge=1, le=1000)

class ConditionalConfig(BaseModel):
    condition_prop: BaseConditionalProps = Field(discriminator='type')

调整后的输入示例：

>>> channel = {"id": "1", "customer_id": "c1", "conditional_config" : {"condition_prop":{"type": "EXPRESSION", "conditional_expression":"ALPHA"}}}
>>> channel_obj = SendNotificationChannel(**channel)
>>> channel_obj.conditional_config.condition_prop
ConditionalExpressionProps(type=<ConditionType.EXPRESSION: 'EXPRESSION'>, conditional_expression='ALPHA')

方案3：Pydantic v1 字段验证器（兼容旧版本） #

如果使用Pydantic v1，用@validator实现类型映射：

from pydantic import validator

class ConditionalConfig(BaseModel):
    cond_type: ConditionType = ConditionType.EXPRESSION
    condition_prop: BaseConditionalProps

    @validator('condition_prop', pre=True, always=True)
    def validate_condition_prop(cls, v, values):
        cond_type = values.get('cond_type')
        if cond_type == ConditionType.EXPRESSION:
            return ConditionalExpressionProps(**v)
        elif cond_type == ConditionType.CYCLE_DUR_TREND:
            return CycleDurationTrendProps(**v)
        return BaseConditionalProps(**v)

测试效果与方案1一致。

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