如何为多个不同的对象生成 Pydantic 模型

Question

我需要一个变量covars，其中包含未知数量的条目，其中每个条目都是三个不同的自定义Pydantic 模型之一。在这种情况下，每个条目都为我的应用程序描述了一个变量。

具体来说，我希望covars 具有以下形式。这里显示了三个条目，即variable1、variable2 和variable3，代表三种不同类型的条目。但是，在部署时，应用程序必须允许接收三个以上的条目，并且并非所有条目类型都需要出现在请求中。

covars = {
            'variable1':  # type: integer
                {
                    'guess': 1,
                    'min': 0,
                    'max': 2,
                },
            'variable2':  # type: continuous
                {
                    'guess': 12.2,
                    'min': -3.4,
                    'max': 30.8,
                },
            'variable3':  # type: categorical
                {
                    'guess': 'red',
                    'options': {'red', 'blue', 'green'},
                }
        }

我已经成功地将三种不同的条目类型创建为三个独立的Pydantic 模型

import pydantic
from typing import Set, Dict, Union


class IntVariable(pydantic.BaseModel):
    guess: int
    min: int
    max: int


class ContVariable(pydantic.BaseModel):
    guess: float
    min: float
    max: float


class CatVariable(pydantic.BaseModel):
    guess: str
    options: Set[str] = {}

注意IntVariable 和ContVariable 之间的数据类型差异。

我的问题：如何制作一个Pydantic 模型，该模型允许组合任意数量的IntVariable、ContVariable 和CatVariable 类型的条目以获得我正在寻找的输出？

计划是使用此模型在数据发布到 API 时验证数据，然后将序列化版本存储到应用程序数据库（使用ormar）。

Answer 1

首先，由于您似乎没有使用预定义的键，您可以使用custom root type，它允许您在 pydantic 模型中拥有任意键名，如 here 所讨论的。接下来，您可以使用Union，它允许模型属性接受不同的类型（并且在定义时也忽略顺序）。因此，无论顺序如何，您都可以传递三个模型的多个条目。

由于IntVariable 和ContVariable 模型具有完全相同数量的属性和键名，因此当将float 数字传递给min 和max 时，它们将转换为int，因为pydantic 没有办法来区分这两个模型。最重要的是，min 和 max 是 Python 中的保留关键字；因此，最好更改它们，如下所示。

from typing import Dict, Set, Union
from pydantic import BaseModel

app = FastAPI()

class IntVariable(BaseModel):
    guess: int
    i_min: int
    i_max: int

class ContVariable(BaseModel):
    guess: float
    f_min: float
    f_max: float


class CatVariable(BaseModel):
    guess: str
    options: Set[str]
    
class Item(BaseModel):
    __root__: Union [IntVariable, ContVariable, CatVariable]

@app.post("/upload")
async def upload(covars: Dict[str, Item]):
    return covars

输入样本如下所示。确保在输入 options Set 时使用方括号 []，否则如果使用大括号 {}，FastAPI 会报错。

{
   "variable1":{
      "guess":1,
      "i_min":0,
      "i_max":2
   },
   "variable2":{
      "guess":"orange",
      "options":["orange", "yellow", "brown"]
   },
   "variable3":{
      "guess":12.2,
      "f_min":-3.4,
      "f_max":30.8
   },
   "variable4":{
      "guess":"red",
      "options":["red", "blue", "green"]
   },
   "variable5":{
      "guess":2.15,
      "f_min":-1.75,
      "f_max":11.8
   }
}

更新

由于上述情况，当为其中一个模型引发ValidationError 时，会引发所有三个模型的错误（而不是仅针对该特定模型引发错误），因此可以使用Discriminated Unions，如this answer 中所述.与歧视性工会，“如果失败，只会引发一个显式错误”.下面的例子：

应用程序.py

from fastapi import FastAPI
from typing import Dict, Set, Union
from pydantic import BaseModel, Field
from typing import Literal

app = FastAPI()

class IntVariable(BaseModel):
    model_type: Literal['int']
    guess: int
    i_min: int
    i_max: int

class ContVariable(BaseModel):
    model_type: Literal['cont']
    guess: float
    f_min: float
    f_max: float


class CatVariable(BaseModel):
    model_type: Literal['cat']
    guess: str
    options: Set[str]
    
class Item(BaseModel):
    __root__: Union[IntVariable, ContVariable, CatVariable] = Field(..., discriminator='model_type')

@app.post("/upload")
async def upload(covars: Dict[str, Item]):
    return covars

测试数据

{
   "variable1":{
      "model_type": "int",
      "guess":1,
      "i_min":0,
      "i_max":2
   },
   "variable2":{
      "model_type": "cat",
      "guess":"orange",
      "options":["orange", "yellow", "brown"]
   },
   "variable3":{
      "model_type": "cont",
      "guess":12.2,
      "f_min":-3.4,
      "f_max":30.8
   },
   "variable4":{
      "model_type": "cat",
      "guess":"red",
      "options":["red", "blue", "green"]
   },
   "variable5":{
      "model_type": "cont",
      "guess":2.15,
      "f_min":-1.75,
      "f_max":11.8
   }
}

另一种解决方案是使用dependency 函数，在其中迭代字典并尝试使用try-catch 块中的三个模型解析字典中的每个项目/条目，类似于this answer (Update 1) 中描述的内容。但是，这将需要遍历所有模型，或者在条目中包含一个鉴别器（例如上面的"model_type"），指示您应该尝试解析哪个模型。

Answer 2

我最终使用自定义验证器解决了这个问题。在此处添加它以补充@Chris 的解决方案。

我使用了其他几个功能来完成这项工作。首先，我将这三种类型设置为Enum 来约束选项。其次，我使用StrictInt、StrictFloat 和StrictStr 来规避python 将int 转换为float 的挑战，如果第一个选项出现在例如guess 是 float，即如果我使用 guess: Union[float,int,str]。第三，我删除输入vtype（类型为VarType）并将其替换为类型为str 的另一个字段type，通过root_validator 使用自定义替换。

import ormar
import pydantic
from enum import Enum
from pydantic import Json, validator, root_validator, StrictInt, StrictFloat, StrictStr
from typing import Set, Dict, Union, Optional
import uuid


class VarType(Enum):
    int = "int"
    cont = "cont"
    cat = "cat"


class Variable(pydantic.BaseModel):
    vtype: VarType
    guess: Union[StrictFloat, StrictInt, StrictStr]
    min: Optional[Union[StrictFloat, StrictInt]] = None
    max: Optional[Union[StrictFloat, StrictInt]] = None
    options: Optional[Set[str]] = None

    # this check is needed to make 'type' available for 'check_guess' validator, but it is not otherwise needed since
    # VarType itself ensures type validation
    @validator('vtype', allow_reuse=True)
    def req_check(cls, t):
        assert t.value in ['int', 'cont', 'cat'], "'vtype' must take value from set ['int', 'cont', 'cat']"
        return t

    # add new field called "type"
    @root_validator(pre=False, allow_reuse=True)
    def insert_type(cls, values):
        if values['vtype'].value == 'int':
            values['type'] = 'int'
        elif values['vtype'].value == 'cont':
            values['type'] = 'float'
        elif values['vtype'].value == 'cat':
            values['type'] = 'str'
        return values

    @root_validator(pre=True, allow_reuse=True)
    def set_guessminmax_types(cls, values):
        if values['vtype'] == 'int':
            values['guess'] = int(values['guess'])
            values['min'] = int(values['min'])
            values['max'] = int(values['max'])
        elif values['vtype'] == 'cont':
            values['guess'] = float(values['guess'])
            values['min'] = float(values['min'])
            values['max'] = float(values['max'])
        return values


    # check right data type of 'guess'
    @validator('guess', allow_reuse=True)
    def check_guess_datatype(cls, g, values):
        if values['vtype'].value == 'int':
            assert isinstance(g, int), "data type mismatch between 'guess' and 'vtype'. Expected type 'int' from 'guess' but received " + str(
                type(g))
            return g
        elif values['vtype'].value == 'cont':
            assert isinstance(g, float), "data type mismatch between 'guess' and 'vtype'. Expected type 'float' from 'guess' but received " + str(
                type(g))
            return g
        elif values['vtype'].value == 'cat':
            assert isinstance(g, str), "data type mismatch between 'guess' and 'vtype'. Expected type 'str' from 'guess' but received " + str(
                type(g))
            return g

    # check that 'min' is included for types 'int', 'cont'
    @validator('min', allow_reuse=True)
    def check_min_included(cls, m, values):
        if values['vtype'].value in ['int', 'cont']:
            assert m is not None
            return m

    # check that 'max' is included for types 'int', 'cont'
    @validator('max', allow_reuse=True)
    def check_max_included(cls, m, values):
        if values['vtype'].value in ['int', 'cont']:
            assert m is not None
            return m

    # check that 'options' is included for type 'cat'
    @validator('options', allow_reuse=True)
    def check_options_included(cls, op, values):
        if values['vtype'].value == 'cat':
            assert op is not None
            return op

    # removes all fields which have value None
    @root_validator(pre=False, allow_reuse=True)
    def remove_all_nones(cls, values):
        values = {k: v for k, v in values.items() if v is not None}
        return values

    class Config:
        fields = {"vtype": {"exclude": True}}