Python 类型提示基础：TypeVar 与联合类型
在 python 中，typing 模块提供了强大的工具来增强代码的可读性和可维护性，其中 typevar 和联合类型（union type）是两个核心概念。
- TypeVar (类型变量)：用于定义泛型函数或类，允许在函数或方法签名中捕获并重用特定类型。例如，一个函数可以接受任意类型的输入并返回相同类型的输出。
- 联合类型 (Union Type)：表示一个值可以是多种指定类型中的任意一种。例如，float | int 表示一个值可以是浮点数或整数。
当这两种类型提示机制结合使用时，可能会遇到一些不直观的类型检查行为，特别是当 TypeVar 带有约束条件时。
约束型 TypeVar 与联合类型的冲突
考虑一个常见的场景：我们希望定义一个泛型函数 f，它能处理 float 或 np.ndarray (或 Fraction) 类型的输入，并返回相同类型的结果。我们可能会自然地使用一个约束型的 TypeVar：
from typing import TypeVar
import numpy as np
from fractions import Fraction
# 示例 1: 使用 numpy.ndarray
T_ndarray = TypeVar("T_ndarray", float, np.ndarray)
def f_ndarray(x: T_ndarray) -> T_ndarray:
"""
期望输入 float 或 np.ndarray，并返回相同类型
"""
return x * 2
# 示例 2: 使用 fractions.Fraction
T_fraction = TypeVar("T_fraction", float, Fraction)
def f_fraction(x: T_fraction) -> T_fraction:
"""
期望输入 float 或 Fraction，并返回相同类型
"""
return x * 2
# 测试调用
print(f_ndarray(1.0))
print(f_ndarray(np.array([1, 2, 3])))
print(f_fraction(1.0))
print(f_fraction(Fraction(1, 2)))现在，假设我们有另一个函数 g，它的输入类型是 float | np.ndarray (或 float | Fraction)，并且在 g 内部调用了 f_ndarray (或 f_fraction)：
# 延续上面的定义
def g_ndarray(x: float | np.ndarray) -> float | np.ndarray:
"""
期望输入 float 或 np.ndarray
"""
return f_ndarray(x) / 2 # Pyright 报错
def g_fraction(x: float | Fraction) -> float | Fraction:
"""
期望输入 float 或 Fraction
"""
return f_fraction(x) / 2 # Pyright 报错在这种情况下，Pyright (或 Mypy 等类型检查器) 会报告一个类型错误，例如：
立即学习“Python免费学习笔记（深入）”；
Argument of type "float | ndarray[Unknown, Unknown]" cannot be assigned to parameter "x" of type "T@f_ndarray" in function "f_ndarray" Type "float | ndarray[Unknown, Unknown]" is incompatible with constrained type variable "T_ndarray"
这个报错令人困惑，因为直观上 x 的类型 (float | np.ndarray) 似乎与 f_ndarray 所期望的类型 (float 或 np.ndarray) 是兼容的。
为什么会报错？
问题的核心在于 TypeVar("T", A, B) 这种约束方式的语义。它表示 T 在调用时必须是 确切的 A 类型，或者 确切的 B 类型。类型检查器在调用 f(x) 时，需要确定 x 的具体类型是 A 还是 B。
然而，当 x 的类型被声明为 A | B (即联合类型) 时，类型检查器知道 x 可能 是 A，也 可能 是 B，但它无法在编译时确定 x 的 单一具体类型。因此，A | B 作为一个整体，不被视为 A 也不被视为 B，所以它与约束型 TypeVar T 不兼容。
为了更好地理解，我们可以简化这个例子：
from typing import TypeVar, Union
from fractions import Fraction
T = TypeVar("T", float, Fraction)
def f(x: T) -> T:
pass
def get_float_or_fraction() -> float | Fraction:
# 模拟一个返回联合类型的函数
return 1.0 # 实际可以是 Fraction(1,2)
num: float | Fraction = get_float_or_fraction()
f(num) # 报错：Argument of type "float | Fraction" cannot be assigned to parameter "x" of type "T@f"这明确指出，float | Fraction 作为一个整体的类型，不能直接赋值给一个被约束为 float 或 Fraction 的 TypeVar。
解决方案
有两种主要的方法可以解决这个问题，具体取决于你的泛型函数 f 的实际需求。
方案一：扩展 TypeVar 约束以包含联合类型
如果你的泛型函数 f 确实需要能够处理 float、np.ndarray 以及 float | np.ndarray 这种联合类型，并且希望当输入是联合类型时，返回类型也保持为该联合类型，那么你需要将联合类型本身添加到 TypeVar 的约束中。
from typing import TypeVar, Union
import numpy as np
from fractions import Fraction
# 方案一示例：扩展 TypeVar 约束
# 针对 numpy.ndarray
T_ndarray_ext = TypeVar("T_ndarray_ext", float, np.ndarray, Union[float, np.ndarray])
def f_ndarray_ext(x: T_ndarray_ext) -> T_ndarray_ext:
"""
现在可以接受 float, np.ndarray 或 float | np.ndarray
"""
return x * 2
def g_ndarray_fixed(x: float | np.ndarray) -> float | np.ndarray:
return f_ndarray_ext(x) / 2 # Pyright 不再报错
# 针对 fractions.Fraction
T_fraction_ext = TypeVar("T_fraction_ext", float, Fraction, Union[float, Fraction])
def f_fraction_ext(x: T_fraction_ext) -> T_fraction_ext:
"""
现在可以接受 float, Fraction 或 float | Fraction
"""
return x * 2
def g_fraction_fixed(x: float | Fraction) -> float | Fraction:
return f_fraction_ext(x) / 2 # Pyright 不再报错
# 示例调用
print(g_ndarray_fixed(1.0))
print(g_ndarray_fixed(np.array([4, 5])))
print(g_fraction_fixed(Fraction(3, 4)))注意事项：
- 这种方法明确告诉类型检查器，T 可能是 float，可能是 np.ndarray，也可能就是 float | np.ndarray 这个联合类型本身。
- 当输入是 float | np.ndarray 时，函数的返回类型也会被推断为 float | np.ndarray。
- 这种方式适用于你希望严格控制 TypeVar 可能的类型范围，并且需要将联合类型作为一个独立的“类型选项”来处理的场景。
方案二：使用 bound 参数定义上限
如果你的泛型函数 f 的目标是接受任何类型，只要它是一个 float 或 np.ndarray 的子类型即可，并且你希望函数返回的类型能尽可能地保留输入的具体类型，那么使用 bound 参数会是更简洁和灵活的选择。
bound 参数定义了一个类型变量的上限，这意味着 TypeVar 可以是这个上限类型或其任何子类型。
from typing import TypeVar, Union
import numpy as np
from fractions import Fraction
# 方案二示例：使用 bound 参数
# 针对 numpy.ndarray
T_ndarray_bound = TypeVar("T_ndarray_bound", bound=Union[float, np.ndarray])
def f_ndarray_bound(x: T_ndarray_bound) -> T_ndarray_bound:
"""
接受任何 float 或 np.ndarray 的子类型
"""
return x * 2
def g_ndarray_bound_fixed(x: float | np.ndarray) -> float | np.ndarray:
return f_ndarray_bound(x) / 2 # Pyright 不再报错
# 针对 fractions.Fraction
T_fraction_bound = TypeVar("T_fraction_bound", bound=Union[float, Fraction])
def f_fraction_bound(x: T_fraction_bound) -> T_fraction_bound:
"""
接受任何 float 或 Fraction 的子类型
"""
return x * 2
def g_fraction_bound_fixed(x: float | Fraction) -> float | Fraction:
return f_fraction_bound(x) / 2 # Pyright 不再报错
# 示例调用
print(g_ndarray_bound_fixed(1.0))
print(g_ndarray_bound_fixed(np.array([7, 8])))
class MyFloat(float):
pass
def get_my_float_or_fraction() -> MyFloat | Fraction:
return MyFloat(1.5)
# 使用 bound 时，返回类型会保留 MyFloat | Fraction
# reveal_type(f_fraction_bound(get_my_float_or_fraction())) # MyFloat | Fraction注意事项：
- bound=Union[A, B] 意味着 T 可以是 A、B，也可以是 A 的子类型，B 的子类型，或者 Union[A, B] 本身。
- 当输入是 float | np.ndarray 时，T 会被推断为 float | np.ndarray。
- 当输入是 MyFloat (一个 float 的子类) 时，T 会被推断为 MyFloat，并且函数返回 MyFloat。这比方案一更灵活，因为它允许更具体的子类型通过。
- 这种方法适用于你希望泛型函数能够处理更广泛的相关类型，并且希望在可能的情况下保留输入类型的具体性的场景。
int 与 float 的特殊情况
在原始问题中，提到了 TypeVar("T", float, int) 在处理 float | int 时不会报错。这可能是因为 int 在 Python 的数值体系中可以无缝地转换为 float，并且类型检查器对 int 和 float 这种常见的数值类型组合有特殊的处理逻辑。在许多情况下，int 可以被视为 float 的一种特殊形式或子类型（在某些操作上下文）。然而，对于 float 和 Fraction 或 float 和 np.ndarray 这种没有直接继承或隐式转换关系的类型，这种特殊处理则不适用，从而暴露了 TypeVar 约束的严格性。
总结
当在 Python 中使用 TypeVar 定义泛型函数，并希望它能接受联合类型（如 float | SomeOtherType）作为输入时，需要注意 TypeVar 的约束方式：
- 约束型 TypeVar("T", A, B)：它期望 T 严格为 A 或 B。因此，A | B 作为一个整体类型，与这种约束不兼容。
- 解决方案一（扩展约束）：如果需要 T 能够精确匹配 A、B 或 A | B，则将 Union[A, B] 明确添加到 TypeVar 的约束列表中：TypeVar("T", A, B, Union[A, B])。
- 解决方案二（使用 bound）：如果希望 T 能接受任何 A | B 的子类型，并且尽可能保留输入类型的具体性，则使用 bound 参数：TypeVar("T", bound=Union[A, B])。
理解这两种方案及其适用场景，能够帮助我们编写出既类型安全又灵活的 Python 泛型代码。在实际开发中，根据函数对输入类型具体性的要求，选择合适的 TypeVar 定义方式至关重要。
