how to define a pydantic model with a field that is a list of a mix of several submodels

Question

I am defining a large model with pydantic and I am stuck trying to define a list where construct multiple submodels, lets see the example:

from pydantic import BaseModel
from typing import Literal, Union

class Model_A(BaseModel):
   type: Literal["A"] = "A"
   name: str
   size: float


class Model_B(BaseModel):
   type: Literal["B"] = "B"
   name: str
   size: float
   color: str
   value: float

class GeneralModel(BaseModel):
   name: str
   objects_list: list[Union[Model_A, Model_B]]


if __name__ == "__main__":
  
  user_list = {
    "name": "John Doe",
    "objects_list": [
        {"name": "box 1", "size": 3},
        {"name": "box 2", "size": 4, "color": "red", "value": 5},
        {"name": "box 3", "size": 6},
    ],
}
  # I would like to obatain the constuction of each model, but my code only construct objects of the first type defined in the Union statement.
  order = GeneralModel(**user_list)
  assert isinstance(order.objects_list[0], Model_A)
  assert isinstance(order.objects_list[1], Model_B)   # Fails, construct Model_A instance

There is any way to solve this and let pydantic to chose the correct model type to use?

Thank you all in advance!

Answer 1

By default, pydantic allows extra fields when building an object. This means that all your objects can be interpreted as Model_A instances, some having extra fields such as color and value. You can prevent this behaviour by adding a configuration to your model:

class Model_A(BaseModel):
    model_config = ConfigDict(extra="forbid")

This means that your first and third objects cannot be interpreted as Model_A anymore, and will be treated as Model_B by default.

You can also use discriminated unions as ndclt suggested, but it is a bit trickier.

Edit: I'm assuming you're using pydantic v2.