如何让Mypy在断言a、b同类型后，通过a的类型收窄推断b为int？

你遇到的问题是：当两个变量a和b的类型都是int | str，且通过断言确保二者类型一致后，收窄a的类型时，Mypy无法自动推断b的类型也同步收窄。以下是几种可行的解决方案：

方案1：自定义类型守卫函数（Type Guard） #

通过编写自定义的类型守卫函数，明确告诉Mypy：当函数返回True时，a和b同时为int或同时为str。这种方式能让Mypy识别两个变量的类型关联。

from typing import TypeGuard

def both_int(a: int | str, b: int | str) -> TypeGuard[int, int]:
    return isinstance(a, int) and isinstance(b, int)

def both_str(a: int | str, b: int | str) -> TypeGuard[str, str]:
    return isinstance(a, str) and isinstance(b, str)

def my_func(a: int | str, b: int | str):
    # 断言二者同为int或同为str
    assert both_int(a, b) or both_str(a, b)
    
    if isinstance(a, int):
        reveal_type(b)  # Mypy现在会正确推断b为int类型

方案2：将变量绑定为联合类型元组 #

如果允许调整函数参数的结构，可以把a和b包装成一个元组，标注为tuple[int, int] | tuple[str, str]。Mypy会自动识别元组内元素的类型绑定关系，收窄其中一个元素的类型时，另一个也会同步收窄。

def my_func(pair: tuple[int, int] | tuple[str, str]):
    a, b = pair
    
    if isinstance(a, int):
        reveal_type(b)  # Mypy推断b为int类型

方案3：分支内补充类型检查（冗余但直接） #

如果不想修改参数结构或新增函数，也可以在if分支内额外检查b的类型。虽然这种写法重复了断言的逻辑，但能让Mypy立即收窄b的类型。

def my_func(a: int | str, b: int | str):
    assert isinstance(a, int) == isinstance(b, int)
    
    if isinstance(a, int) and isinstance(b, int):
        reveal_type(b)  # Mypy推断b为int类型

总结 #

最推荐的是自定义类型守卫函数，既保持代码的简洁性和可读性，又能让Mypy准确推断类型关联。如果参数结构允许，使用联合类型元组的方式会更简洁直观。

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