如何让mypy识别类变量var_b的类型依赖于var_a的取值？

给定以下Python代码：

class A:
    var_a: bool = False
    var_b: int | str  # str if var_a is True, else int


a = A()
a.var_a = True
a.var_b # should be str

需要为这段代码添加合适的类型标注，使得mypy能够识别：当var_a为True时，var_b的类型应为str；当var_a为False时，var_b的类型应为int。请问该如何实现？是否可以通过灵活运用Literal[True]或Literal[False]来达成这一目标？

可以通过泛型 + Literal类型实现这种属性间的类型依赖，普通的bool类型无法让mypy跟踪var_a的具体值对var_b类型的约束，具体步骤如下：

• 导入typing模块中的Generic、TypeVar、Literal工具
• 定义绑定Literal[True]和Literal[False]的TypeVar
• 让类A继承Generic并传入该TypeVar作为泛型参数
• 根据泛型参数的取值，分别约束var_a和var_b的类型

具体代码实现：

from typing import Generic, TypeVar, Literal

# 限定只能取Literal[True]或Literal[False]的TypeVar
BoolType = TypeVar("BoolType", Literal[True], Literal[False])

class A(Generic[BoolType]):
    var_a: BoolType
    var_b: str if BoolType is Literal[True] else int

    def __init__(self, var_a: BoolType = False) -> None:
        self.var_a = var_a
        # 初始化时匹配对应类型的默认值
        self.var_b = "" if var_a else 0

# 使用示例
a_true: A[Literal[True]] = A(True)
a_true.var_b = "hello"  # mypy识别为str，无报错
a_true.var_b = 123      # mypy抛出类型错误

a_false: A[Literal[False]] = A(False)
a_false.var_b = 456     # mypy识别为int，无报错
a_false.var_b = "test"  # mypy抛出类型错误

如果需要支持后续修改var_a并让mypy自动更新var_b的类型，需要通过专门的方法来处理（直接修改属性时mypy无法自动切换泛型参数），示例如下：

from typing import cast

class A(Generic[BoolType]):
    # 省略上述__init__等代码...

    def set_var_a(self, new_val: Literal[True] | Literal[False]) -> None:
        if new_val is True:
            # 强制转换泛型参数和属性类型
            self.var_a = cast(Literal[True], new_val)
            self.var_b = cast(str, "")
        else:
            self.var_a = cast(Literal[False], new_val)
            self.var_b = cast(int, 0)

注意：这种实现依赖mypy对泛型和Literal类型的支持，请确保使用mypy 0.910及以上版本。

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