如何為多個不同的對象生成 Pydantic model

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 model 允許組合任意數量的IntVariable 、 ContVariable和CatVariable類型的條目以獲得我正在尋找的 output ？

計划是使用這個 model 來驗證發送到 API 的數據，然后將序列化版本存儲到應用程序數據庫（使用ormar ）。

Answer 1

首先，由於您似乎沒有使用預定義的鍵，您可以使用自定義根類型，它允許您在 pydantic model 中使用任意鍵名，如此處討論。 接下來，您可以使用Union ，它允許 model 屬性接受不同的類型（並且在定義時也忽略順序）。 因此，無論順序如何，您都可以傳遞三個模型的多個條目。

由於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 ，如本答案中所述。 對於區分聯合， “在失敗的情況下只會引發一個顯式錯誤” 。 示例如下：

應用程序.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
   }
}

另一種解決方案是擁有依賴項 function，在其中迭代字典並嘗試使用 try-catch 塊中的三個模型解析字典中的每個項目/條目，類似於此答案（更新 1）中描述的內容。 但是，這將需要遍歷所有模型，或者在條目中有一個鑒別器（例如上面的"model_type" ），指示您應該嘗試解析哪個 model。

Answer 2

我最終使用自定義驗證器解決了這個問題。 在此處添加它以補充@Chris 的解決方案。

我使用了一些其他功能來完成這項工作。 首先，我將這三種類型設置為Enum來約束選項。 其次，我使用StrictInt 、 StrictFloat和StrictStr來規避挑戰，如果guess中出現的第一個選項是float ，即python會將int轉換為float ，即如果我要使用guess: Union[float,int,str] 。 第三，我刪除了輸入vtype （類型VarType ），並通過root_validator使用自定義替換將其替換為另一個類型為str的字段type 。

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}}