如何基于Union创建约束型TypeVar？解决mypy类型检查报错

正常通过类型检查的代码 #

当直接在TypeVar中枚举类型时，mypy可以正确推导类型：

from typing import TypeVar

S = TypeVar("S", int, float, complex)

def func(x: list[S], m: S) -> list[S]:
    return [val * m for val in x]

out1: list[int] = func([1, 2, 3], 4)
out2: list[complex] = func([1., 2., 3.], 4.)

使用外部Union+bound约束时的错误场景 #

如果Number是来自外部包无法修改的联合类型，使用bound约束TypeVar会触发mypy错误：

from typing import TypeVar

Number = int | float | complex  # 实际为外部包定义，无法修改
S = TypeVar("S", bound=Number)

def func(x: list[S], m: S) -> list[S]:
    return [val * m for val in x]

out1: list[int] = func([1, 2, 3], 4)
out2: list[complex] = func([1., 2., 3.], 4.)

mypy报错信息：

y.py:7: error: List comprehension has incompatible type List[Union[int, float, complex]]; expected List[S]  [misc]
y.py:7: error: Unsupported operand types for * (likely involving Union)  [operator]

解决方法 #

方法1：提取外部Union的成员作为TypeVar的显式约束 #

利用typing.get_args获取外部Union类型的成员，将其作为TypeVar的枚举参数，效果和直接枚举类型一致：

from typing import TypeVar, get_args
from external_package import Number  # 导入外部包的联合类型

# 提取Number的所有成员类型作为TypeVar的约束
S = TypeVar("S", *get_args(Number))

def func(x: list[S], m: S) -> list[S]:
    return [val * m for val in x]

out1: list[int] = func([1, 2, 3], 4)
out2: list[complex] = func([1., 2., 3.], 4.)

方法2：使用Protocol定义可乘法操作的类型 #

如果外部Union的成员都支持乘法操作，可以定义一个协议来约束类型，避免Union带来的推导问题：

from typing import TypeVar, Protocol
from external_package import Number

class Multipliable(Protocol):
    def __mul__(self, other: "Multipliable") -> "Multipliable": ...

S = TypeVar("S", bound=Multipliable)

def func(x: list[S], m: S) -> list[S]:
    return [val * m for val in x]

out1: list[int] = func([1, 2, 3], 4)
out2: list[complex] = func([1., 2., 3.], 4.)

注意：此方法需确保外部Number的所有成员都符合Multipliable协议的定义，mypy会自动验证这一点。

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