Mypy误判返回Union[A,B]的工厂函数返回Parent的原因及优化方案

复现代码 #

from typing import Union


class Parent:
    def __init__(self) -> None:
        print("In Parent")


class A(Parent):
    def __init__(self) -> None:
        super().__init__()
        print("In A")


class B(Parent):
    def __init__(self) -> None:
        super().__init__()
        print("In B")


def factory(arg: str) -> Union[A, B]:
    mapper = {"A": A, "B": B}

    return mapper[arg]()

Mypy报错信息 #

error: Incompatible return value type (got "Parent", expected "Union[A, B]")

疑问 #

这段代码要么在参数非法时报错，要么返回Union中指定的子类实例，为何Mypy会认为返回的是Parent类型？是否有更优的写法？以下分支判断的变体没有报错，但子类较多时更倾向于字典映射的写法：

def factory(arg: str) -> Union[A, B]:
    if arg == "A":
        return A()

    return B()

报错原因 #

Mypy对字典类型的推导是保守的：当定义mapper = {"A": A, "B": B}时，它会自动将字典值的类型推导为type[Parent]——因为A和B都是Parent的子类，Mypy默认取它们的公共父类作为字典值的统一类型。因此调用mapper[arg]()时，Mypy认为返回的是Parent实例，和函数标注的Union[A, B]返回类型冲突。

优化写法 #

以下几种方式既能保留字典映射的简洁性，又能让Mypy正确识别类型：

1. 显式标注字典类型 #

直接给mapper指定精确的类型，告诉Mypy字典的值是A或B的类：

from typing import Union, Dict, Type

# 类定义不变

def factory(arg: str) -> Union[A, B]:
    mapper: Dict[str, Union[Type[A], Type[B]]] = {"A": A, "B": B}
    return mapper[arg]()

2. 用Literal限定输入参数（推荐） #

如果factory的入参arg只能是"A"或"B"，用Literal明确输入范围后，Mypy能精准推导返回类型：

from typing import Union, Literal

# 类定义不变

def factory(arg: Literal["A", "B"]) -> Union[A, B]:
    mapper = {"A": A, "B": B}
    return mapper[arg]()

3. 用cast强制类型转换 #

如果不想修改字典类型标注，也可以用cast告知Mypy实际返回类型，但这种方式属于绕过类型检查，仅作为临时方案：

from typing import Union, cast

# 类定义不变

def factory(arg: str) -> Union[A, B]:
    mapper = {"A": A, "B": B}
    return cast(Union[A, B], mapper[arg]())

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