如何检查TypeVar实际类型并消除mypy类型不兼容错误

问题背景 #

定义了受限类型变量T，并实现了env函数，接收T的类型对象作为参数，返回对应类型的值或None。代码运行正常，但mypy报出返回值类型不兼容的错误。

mypy错误信息 #

test_generic.py:20: error: Incompatible return value type (got "bool", expected "str | None")  [return-value]
test_generic.py:22: error: Incompatible return value type (got "int", expected "bool | None")  [return-value]
test_generic.py:22: error: Incompatible return value type (got "int", expected "str | None")  [return-value]
test_generic.py:24: error: Incompatible return value type (got "str", expected "bool | None")  [return-value]
test_generic.py:24: error: Incompatible return value type (got "str", expected "int | None")  [return-value]

原代码 #

from typing import TypeVar, Type, Optional

T = TypeVar("T", bool, int, str)


def get_str() -> str:
    return "str"


def get_int() -> int:
    return 1


def get_bool() -> bool:
    return True


def env(value_type: Type[T]) -> Optional[T]:
    if value_type is bool:
        return get_bool()
    elif value_type is int:
        return get_int()
    elif value_type is str:
        return get_str()
    else:
        return None


def test_int() -> None:
    assert env(int) == get_int()


def test_bool() -> None:
    assert env(bool) == get_bool()


def test_str() -> None:
    assert env(str) == get_str()

使用版本 #

$ python --version
Python 3.11.1
$ mypy --version
mypy 1.5.1 (compiled: yes)

错误原因 #

mypy无法自动关联value_type is bool这类类型检查与TypeVar T的绑定关系。当你在分支中判断value_type是具体类型时，mypy不会推断出此时T应该就是该具体类型，而是认为返回值需要匹配T的所有可能类型（bool/int/str），因此抛出类型不兼容的错误。

修复方案 #

使用**函数重载（Overload）**明确声明不同输入类型对应的返回类型，让mypy能正确识别每个分支的类型匹配关系。

修复后的代码 #

from typing import TypeVar, Type, Optional, overload

T = TypeVar("T", bool, int, str)


def get_str() -> str:
    return "str"


def get_int() -> int:
    return 1


def get_bool() -> bool:
    return True


@overload
def env(value_type: Type[bool]) -> Optional[bool]: ...
@overload
def env(value_type: Type[int]) -> Optional[int]: ...
@overload
def env(value_type: Type[str]) -> Optional[str]: ...
def env(value_type: Type[T]) -> Optional[T]:
    if value_type is bool:
        return get_bool()
    elif value_type is int:
        return get_int()
    elif value_type is str:
        return get_str()
    else:
        return None


def test_int() -> None:
    assert env(int) == get_int()


def test_bool() -> None:
    assert env(bool) == get_bool()


def test_str() -> None:
    assert env(str) == get_str()

这样mypy就能根据重载声明，正确识别传入不同类型时的返回值类型，消除错误提示。

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