How to return a response with a list of different Pydantic models using FastAPI?

Question

Context

I am creating an API using FastAPI to compute the shortest path on a graph. My response consists of nodes and relations.

These nodes can be of different types, which means they can have different attributes:

class SchoolNode(BaseModel):
    uid: int
    code: Optional[str]
    label: str = 'school'


class PersonNode(BaseModel):
    uid: int
    name: Optional[str]
    surname: Optional[str]
    label: str = 'person'


class PetNode(BaseModel):
    uid: int
    name: Optional[str]
    surname: Optional[str]
    label: str = 'pet'

My response follows this format:

class Response(BaseModel):
    links: List[Link]
    nodes: List[Union[SchoolNode, PersonNode, PetNode]]
    start: int
    end: int

Note that I cannot change the response format since my output will be used by a custom library based on d3js that needs data in input in this specific format.

You can see a gist with full code example here .

Problem

The API run successfully, but the response inside 'nodes' property is unable to understand which model must be chosen. The expected output is:

{
    'links': [
        {'source': 1, 'target': 123, 'type': 'GO_TO_SCHOOL'},
        {'source': 100, 'target': 123, 'type': 'GO_TO_SCHOOL'},
    ],
    'nodes': [
        {'uid': 1, 'label': 'person', 'name': 'Bob', 'surname': 'Foo'},
        {'uid': 123, 'label': 'school', 'code': 'ABCD'},
        {'uid': 100, 'label': 'person', 'name': 'Alice', 'surname': 'Bar'}
    ],
    'start': 1,
    'end': 100
}

while the obtained output is:

{
    "links": [
        {"source": 1, "target": 123, "type": "GO_TO_SCHOOL"},
        {"source": 123, "target": 100, "type": "GO_TO_SCHOOL"}
    ],
    "nodes": [
        {"uid": 1, "code": null, "label": "person"},
        {"uid": 123, "code": "ABCD", "label": "school"},
        {"uid": 100, "code": null, "label": "person"}
    ],
    "start": 1,
    "end": 100
}

Here you can see how the first and third nodes show the attributes of the first node (SchoolNode) instead of the correct ones (PersonNode)

Question

How should I change my Response to return the correct output? I tried using an if-then-else logic like

nodes = []
for node in graph['nodes']:
    if node['label'] == 'person':
        node.append(PersonNode(**node)
    elif:
        ...

but nothing changed.

I also tried using Field(..., discriminator='label') and I guess this is the correct way to address this issue but without success at the moment.

Any help is appreciated, thanks in advance!

Answer 1

Thanks to @ Chris and following the links he sent me I can solve this problem.

The solution was creating a unique model UnionNode with a __root__ property with Field(..., discriminator='label') . Moreover, label property in nodes must have a Literal typing.

class SchoolNode(BaseModel):
    id: int
    label: Literal['school']
    code: Optional[str]


class PersonNode(BaseModel):
    id: int
    label: Literal['person']
    name: Optional[str]
    surname: Optional[str]


class PetNode(BaseModel):
    id: int
    label: Literal['pet']
    name: Optional[str]
    surname: Optional[str]


class UnionNode(BaseModel):
    __root__: Union[SchoolNode, PersonNode, PetNode] = Field(..., discriminator='label')


class Response(BaseModel):
    links: List[Link]
    nodes: List[UnionNode]
    start: int
    end: int