如何让Pyright正确推断索引后ndarray的类型？

问题原因 #

Pyright在处理Numpy多维数组的整数索引（如[0]）时，无法自动推导索引后数组的维度变化。尽管batch_str_to_int_array的返回类型明确标注为三维数组，Pyright当前没有内置逻辑将索引操作映射为“移除一个维度、保持元素类型”的类型收缩，导致索引后的变量被推断为Any类型。这是静态类型检查器与Numpy动态数组类型系统适配的局限性问题。

正确实现方式 #

方法1：显式类型标注 #

直接为索引后的变量添加预期的类型注解，明确告知Pyright变量类型：

import numpy as np

def str_to_int_array(foo: str) -> np.ndarray[tuple[int, int], np.dtype[np.int_]]:
    3D_int_array = batch_str_to_int_array([foo])  # 返回np.ndarray[tuple[int, int, int], np.dtype[np.int_]]
    2D_int_array: np.ndarray[tuple[int, int], np.dtype[np.int_]] = 3D_int_array[0]
    return 2D_int_array

方法2：使用typing.cast强制类型转换 #

通过cast工具明确告诉静态检查器索引后的变量类型，适用于需要强调类型转换逻辑的场景：

from typing import cast
import numpy as np

def str_to_int_array(foo: str) -> np.ndarray[tuple[int, int], np.dtype[np.int_]]:
    3D_int_array = batch_str_to_int_array([foo])
    2D_int_array = cast(np.ndarray[tuple[int, int], np.dtype[np.int_]], 3D_int_array[0])
    return 2D_int_array

方法3：简化类型注解（忽略维度严格校验） #

若无需严格标注维度，可使用numpy.typing.NDArray简化写法，仅指定元素类型：

import numpy as np
from numpy.typing import NDArray

def str_to_int_array(foo: str) -> NDArray[np.int_]:
    3D_int_array: NDArray[np.int_] = batch_str_to_int_array([foo])
    2D_int_array: NDArray[np.int_] = 3D_int_array[0]
    return 2D_int_array

补充说明 #

当前Pyright对Numpy多维数组索引的类型推断支持尚未完善，显式标注类型是最直接有效的解决方案，既可以让Pyright正确识别类型，也能提升代码的可读性与维护性。

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