基于枚举参数的条件返回值键入

Question

我有一些pydantic basemodels，我想用值填充。

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"]},
    },
}

为了将值加载到正确的基本模型类中，我创建了一个系统类，该类也需要其他地方并且具有更多功能，但是我省略了详细信息以清晰。

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]

然后，通过系统[“ A”]。文档我可以直接访问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> load_document具有union> Union的键入返回值[docucta，documenta，documenta，documentb，documentb]，显然documenta和documentc无法访问列属性。但是，这里可以加载的唯一文档类型是documentb无论如何。

我知道我可以在处理程序功能之外加载文档并改用它，但是我希望将其加载到处理程序中。

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

是否可以有条件键入基于枚举输入值的函数返回值？

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)

现在代码的其余代码现在（游乐场）。