如何使用 FastAPI 返回包含不同 Pydantic 模型列表的响应？

Question

语境

我正在使用 FastAPI 创建一个 API 来计算图形上的最短路径。 我的回应由节点和关系组成。

这些节点可以是不同的类型，这意味着它们可以具有不同的属性：

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'

我的回复遵循以下格式：

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

请注意，我无法更改响应格式，因为我的 output 将由基于 d3js 的自定义库使用，该库需要以这种特定格式输入数据。

您可以在此处查看包含完整代码示例的要点。

问题

API 成功运行，但“节点”属性内的响应无法理解必须选择哪个 model。 预期的 output 是：

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

而得到的output是：

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

在这里您可以看到第一个和第三个节点如何显示第一个节点 (SchoolNode) 的属性而不是正确的属性 (PersonNode)

问题

我应该如何更改我的响应以返回正确的 output？ 我尝试使用类似 if-then-else 的逻辑

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

但没有任何改变。

我也尝试使用Field(..., discriminator='label')我想这是解决这个问题的正确方法，但目前没有成功。

感谢您的帮助，在此先感谢！

Answer 1

感谢@Chris并点击他发给我的链接，我可以解决这个问题。

解决方案是创建一个独特的 model UnionNode ，其__root__属性带有Field(..., discriminator='label') 。 此外，节点中的 label 属性必须具有Literal类型。

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