使用ArrayLike时触发Mypy错误，求兼容ndarray与列表的正确用法

问题背景 #

当尝试用npt.ArrayLike定义兼容NumPy数组（ndarray）和普通列表的函数签名时，Mypy会抛出类型错误——因为npt.ArrayLike的范围过于宽泛，包含了单个整数、布尔值、字符串等不支持len()或索引操作的类型，导致直接调用这些操作时触发检查失败，且希望不通过cast绕过类型检查。

你的示例代码：

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

def f(a: npt.ArrayLike) -> int:
    return len(a)

def g(a: npt.ArrayLike) -> Any:
    return a[0]

print(f(np.array([0, 1])), g(np.array([0, 1])))
print(f([0, 1]), g([0, 1]))

Mypy报错信息（翻译后）：

参数1传入"len"时类型不兼容，预期类型"Sized"，实际类型为"Union[_SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], bool, int, float, complex, str, bytes, _NestedSequence[Union[bool, int, float, complex, str, bytes]]]"  [arg-type]

类型"Union[_SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], bool, int, float, complex, str, bytes, _NestedSequence[Union[bool, int, float, complex, str, bytes]]]"不可被索引  [index]

解决方案 #

方案1：自定义严格的兼容类型别名 #

创建一个仅包含ndarray和标准序列类型的类型别名，缩小类型范围以匹配实际需求：

import numpy as np
import numpy.typing as npt
from typing import Sequence, TypeVar, Any

# 定义泛型类型变量，支持更灵活的类型标注
T = TypeVar("T")
# 自定义类型：兼容NumPy数组和任意序列（列表、元组等）
ArrayOrSequence = npt.NDArray[T] | Sequence[T]

def f(a: ArrayOrSequence[Any]) -> int:
    return len(a)

def g(a: ArrayOrSequence[Any]) -> Any:
    return a[0]

print(f(np.array([0, 1])), g(np.array([0, 1])))
print(f([0, 1]), g([0, 1]))

方案2：直接使用联合类型 #

如果不需要泛型支持，可直接在函数签名中写联合类型：

import numpy as np
import numpy.typing as npt
from typing import Sequence, Any

def f(a: npt.NDArray[Any] | Sequence[Any]) -> int:
    return len(a)

def g(a: npt.NDArray[Any] | Sequence[Any]) -> Any:
    return a[0]

print(f(np.array([0, 1])), g(np.array([0, 1])))
print(f([0, 1]), g([0, 1]))

两种方案都能让Mypy正确识别参数支持len()和索引操作，同时兼容NumPy数组和普通列表，无需使用cast绕过类型检查。

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