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

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!

Please have a look here and here on how to use Discriminated Unions. — Chris
– Chris, Commented Oct 4, 2022 at 8:45
Thanks for the quick response. Since I have a List[Union[Model1, Model2]], should I have to follow your second link and use the __root__ property? — ndricca
– ndricca, Commented Oct 4, 2022 at 9:00
the two questions are slightly different but probably differences are not relevant, if it could help it's ok to mark this as a duplicate. thanks again! — ndricca
– ndricca, Commented Oct 4, 2022 at 9:21

ndricca · Accepted Answer · 2022-10-04 09:16:30Z

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

answered Oct 4, 2022 at 9:16

ndricca

5424 silver badges13 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

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

Context

Problem

Question

1 Answer 1

Comments

Linked

Hot Network Questions

Collectives™ on Stack Overflow

Context

Problem

Question

1 Answer 1

Comments

Linked

Related