基于枚举参数的条件返回值类型

Question

我有一些 pydantic BaseModel，我想用值填充它们。

from enum import Enum
from typing import Dict, List, Literal, Type, Union, overload    

from pydantic import BaseModel


class Document(BaseModel):
    name: str
    pages: int


class DocumentA(Document):
    reviewer: str


class DocumentB(Document):
    columns: Dict[str, Dict]


class DocumentC(Document):
    reviewer: str
    tools: List[str]

示例值：

db = {
    "A": {
        0: {"name": "Document 1", "pages": 2, "reviewer": "Person A"},
        1: {"name": "Document 2", "pages": 3, "reviewer": "Person B"},
    },
    "B": {
        0: {"name": "Document 1", "pages": 1, "columns": {"colA": "A", "colB": "B"}},
        1: {"name": "Document 2", "pages": 5, "columns": {"colC": "C", "colD": "D"}},
    },
    "C": {
        0: {"name": "Document 1", "pages": 7, "reviewer": "Person C", "tools": ["hammer"]},
        1: {"name": "Document 2", "pages": 2, "reviewer": "Person A", "tools": ["hammer", "chisel"]},
    },
}

为了将值加载到正确的 BaseModel 类中，我创建了一个系统类，它在其他地方也需要并且具有更多功能，但为了清楚起见，我省略了详细信息。

class System(Enum):
    A = ("A", DocumentA)
    B = ("B", DocumentB)
    C = ("C", DocumentC)

    @property
    def key(self)-> str:
        return self.value[0]

    @property
    def Document(self) -> Union[Type[DocumentA], Type[DocumentB], Type[DocumentC]]:
        return self.value[1]

然后，通过

System["A"].Document

我可以直接访问

DocumentA

。 为了加载值，我使用这个函数（暂时忽略处理 IndexErrors）：

def load_document(db: Dict, idx: int, system: System) -> Union[DocumentA, DocumentB, DocumentC]:
    data = db[system.key][idx]
    return system.Document(**data)

现在，我可能需要处理一些 B 类型的数据，我直接在处理函数中加载这些数据。

def handle_document_B(db: Dict, idx: int):
    doc = load_document(db=db, idx=idx, system=System.B)
    # Following line raises mypy errors
    # Item "DocumentA" of "Union[DocumentA, DocumentB, DocumentC]" has no attribute "columns"
    # Item "DocumentC" of "Union[DocumentA, DocumentB, DocumentC]" has no attribute "columns"
    print(doc.columns)

运行 mypy 会在

print(doc.columns)

行上引发错误，因为

load_document

的返回值为

Union[DocumentA, DocumentB, DocumentC]

，并且显然

DocumentA

和

DocumentC

无法访问

columns

属性。但无论如何，这里唯一可以加载的文档类型是

DocumentB

。

我知道我可以在处理程序函数之外加载文档并传递它，但我更愿意将其加载到处理程序中。

我通过使用正确的 Document 类重载

load_document

函数来规避类型问题，但这似乎是一个乏味的解决方案，因为我需要为将来可能添加的每个系统手动添加一个重载器。

是否可以根据 Enum 输入值有条件地键入提示函数返回值？

Answer 1

您可以根据所选选项显式注释返回类型：

from typing import Literal, overload

@overload
def load_document(db: Dict, idx: int, system: Literal[System.A]) -> DocumentA: ...

@overload
def load_document(db: Dict, idx: int, system: Literal[System.B]) -> DocumentB: ...

@overload
def load_document(db: Dict, idx: int, system: Literal[System.C]) -> DocumentC: ...

def load_document(db: Dict, idx: int, system: System) -> Union[DocumentA, DocumentB, DocumentC]:
    data = db[system.key][idx]
    return system.Document(**data)

现在对其余代码进行类型检查（playground）。