How to generate Pydantic model for multiple different objects

Question

I need to have a variable covars that contains an unknown number of entries, where each entry is one of three different custom Pydantic models. In this case, each entry describes a variable for my application.

Specifically, I want covars to have the following form. It is shown here for three entries, namely variable1, variable2 and variable3, representing the three different types of entries. Though, when deployed, the application must allow to receive more than three entries, and not all entry types need to be present in a request.

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

I have successfully created the three different entry types as three separate Pydantic models

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] = {}

Notice the data type difference between IntVariable and ContVariable.

My question: How to make a Pydantic model that allows combining any number of entries of types IntVariable, ContVariable and CatVariable to get the output I am looking for?

The plan is to use this model to verify the data as it is being posted to the API, and then store a serialized version to the application db (using ormar).

Answer 1

First, since you don't seem to be using pre-defined keys, you could use a custom root type, which allows you to have arbitrary key names in a pydantic model, as discussed here. Next, you could use a Union, which allows a model attribute to accept different types (and also ignores the order when defined). Thus, you can pass a number of entries of your three models, regardless of the order.

Since IntVariable and ContVariable models have exactly the same number of attributes and key names, when passing float numbers to min and max, they are converted to int, as there is no way for pydantic to differentiate between the two models. On top of that, min and max are reserved keywords in Python; thus, it would be preferable to change them, as shown below.

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

Input sample is shown below. Make sure to use square brackets [] when typing the options Set, as FastAPI would otherwise complain, if braces {} were used.

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

Update

Since with the above, when a ValidationError is raised for one of the models, errors for all three models are raised (instead of raising errors only for that specific model), one could use Discriminated Unions, as described in this answer. With Discriminated Unions, "only one explicit error is raised in case of failure". Example below:

app.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

Test data

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

An alternative solution would be to have a dependency function, where you iterate over the dictionary and try parsing each item/entry in the dictionary using the three models within a try-catch block, similar to what described in this answer. However, that would require either looping through all the models, or having a discriminator in the entry (such as "model_type" above), indicating which model you should try parsing.