如何基于Enum参数实现Python函数条件返回值类型标注

不需要为每个新增的枚举值手动编写@overload签名，通过泛型绑定枚举成员关联的模型类型即可实现自动类型推断，mypy可精准识别传入不同枚举成员时的返回值类型。

具体实现步骤 #

1. 导入泛型相关的类型工具，定义绑定到Document基类的类型变量，限定泛型范围为Document的子类：

from typing import Generic, TypeVar, Type, Dict
from enum import Enum
from pydantic import BaseModel

DocT = TypeVar("DocT", bound=Document)

2. 改造System枚举，让其继承泛型类Generic[DocT]，为每个枚举成员添加精确的泛型类型注解，同时将Document属性的返回值改为对应泛型的类类型：

class System(Enum, Generic[DocT]):
    # 为每个成员标注对应的泛型参数，明确该成员关联的Document子类
    A: "System[DocumentA]" = ("A", DocumentA)
    B: "System[DocumentB]" = ("B", DocumentB)
    C: "System[DocumentC]" = ("C", DocumentC)

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

    @property
    def Document(self) -> Type[DocT]:
        return self.value[1]

如果使用Python 3.10以下版本，可在文件头部添加from __future__ import annotations，去掉成员注解的引号即可。

3. 改造load_document为泛型函数，将入参system的类型标注为System[DocT]，返回值直接标注为DocT即可：

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

效果验证 #

改造完成后，类型检查器会自动根据传入的枚举成员推断返回值类型，不需要额外的类型断言或重载：

def handle_document_B(db: Dict, idx: int):
    doc = load_document(db=db, idx=idx, system=System.B)
    # mypy可正确推断doc为DocumentB类型，访问columns属性无报错
    print(doc.columns)

def handle_document_A(db: Dict, idx: int):
    doc = load_document(db=db, idx=idx, system=System.A)
    # mypy可正确推断doc为DocumentA类型，访问reviewer属性无报错
    # 如果尝试访问doc.columns会直接抛出类型错误，符合预期
    print(doc.reviewer)

后续扩展 #

新增文档类型时，只需要在System枚举中添加对应的成员注解即可，不需要修改load_document的任何代码，维护成本远低于手动编写重载：

# 新增DocumentD类型时仅需修改枚举类
class DocumentD(Document):
    export_path: str

class System(Enum, Generic[DocT]):
    A: "System[DocumentA]" = ("A", DocumentA)
    B: "System[DocumentB]" = ("B", DocumentB)
    C: "System[DocumentC]" = ("C", DocumentC)
    # 新增一行即可完成适配
    D: "System[DocumentD]" = ("D", DocumentD)
    # ... 其他原有代码保持不变

内容的提问来源于stack exchange，提问作者fstermann