Python 中类型变量和类型联合不兼容吗？

Question

考虑以下代码片段：

from typing import TypeVar

import numpy as np

T = TypeVar("T", float, np.ndarray)


def f(x: T) -> T:
    """
    expects a float or an array and returns an output of the same type
    """
    return x * 2


f(1)  # ok

f(np.array([1, 2, 3]))  # ok


def g(x: float | np.ndarray) -> float | np.ndarray:
    """
    expects either a float or an array
    """
    return f(x) / 2  # should be fine, but pyright complains about type

我创建了一个 TypeVar 来暗示

f

期望输入浮点数或数组，并将返回相同类型的输出。

g

中的类型提示更加宽松。它需要一个浮点或一个数组，并将返回一个浮点或数组，而不将输出的类型限制为输入的类型。

直观上来说，这个设置是有道理的。在

g

函数的定义中，我们知道我们期望

x

是浮点数或数组，即

f

期望作为输入。然而，当我在最后一行将

x

传递到

f

时，Pyright 抱怨道：

“float | ndarray[Unknown, Unknown]”类型的参数不能分配给函数“f”中“T@f”类型的参数“x”

类型“float | ndarray[Unknown, Unknown]”与约束类型变量“T”不兼容

这是令人惊讶和令人沮丧的，因为这意味着如果不非常谨慎地编写类型提示的方式，就无法使用我的函数

f

。

对于如何解决这个问题有什么想法吗？

Answer 1

虽然不是理想的解决方案，但您可以通过在显式类型保护下重复调用

f

中的

g

来解决此错误：

def g(x: float | np.ndarray) -> float | np.ndarray:
    """
    expects either a float or an array
    """
    if isinstance(x, (float, int)):
        return f(x) / 2
    return f(x) / 2

使用pyright v1.1.329输出：

0 errors, 0 warnings, 0 informations