如何使用 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