Pylance中Union类型变量的类型推断问题

我编写了一个根据参数为标量或NumPy数组执行不同逻辑的函数，但遇到两个Pylance类型错误无法解决：

错误信息 #

1. Expression of type "generic | bool | int | float | complex | str | bytes | memoryview" cannot be assigned to return type "float" Type "generic | bool | int | float | complex | str | bytes | memoryview" cannot be assigned to type "float" "bytes" is incompatible with "float" PylancereportGeneralTypeIssues
2. "__getitem__" method not defined on type "float" PylancereportGeneralTypeIssues

示例代码 #

import numpy as np
import numpy.typing as npt


def func(x: float | npt.NDArray[np.float64]) -> float:
    if np.isscalar(x):
        return x     # Error 1)
    else:
        return x[0]  # Error 2)


def main():
    x = 0.4
    print(func(x))


if __name__ == "__main__":
    main()

解决方案 #

问题根源是Pylance无法通过np.isscalar()的结果自动窄化参数类型，需要手动添加**类型守卫（TypeGuard）**来明确分支中的类型。

方法1：自定义TypeGuard函数 #

通过typing.TypeGuard创建一个类型守卫函数，告诉Pylance当判断为True时，参数是float类型；反之则是NumPy数组类型。

修改后的代码：

import numpy as np
import numpy.typing as npt
from typing import TypeGuard


def is_float_scalar(x: float | npt.NDArray[np.float64]) -> TypeGuard[float]:
    # 明确返回True时x是float类型的标量
    return np.isscalar(x)


def func(x: float | npt.NDArray[np.float64]) -> float:
    if is_float_scalar(x):
        return x     # Pylance现在识别x为float，无错误
    else:
        return x[0]  # Pylance自动推断x为NDArray，支持下标访问


def main():
    x = 0.4
    print(func(x))
    arr = np.array([1.5, 2.3])
    print(func(arr))


if __name__ == "__main__":
    main()

方法2：结合isinstance做类型判断 #

如果不想额外定义函数，也可以直接在分支中用isinstance结合np.generic（覆盖NumPy原生标量类型）来辅助类型窄化：

import numpy as np
import numpy.typing as npt


def func(x: float | npt.NDArray[np.float64]) -> float:
    if isinstance(x, (float, np.generic)) and np.isscalar(x):
        return x
    else:
        return x[0]


def main():
    x = 0.4
    print(func(x))
    arr = np.array([1.5, 2.3])
    print(func(arr))


if __name__ == "__main__":
    main()

两种方法都能让Pylance正确识别不同分支的参数类型，解决类型错误问题。

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