如何为兼容list和np.ndarray的Python函数参数添加类型注解？

针对你的需求，有几种实用的方案可以实现同时兼容list和np.ndarray，且保证类型检查工具能识别它们的__len__和__getitem__方法：

1. 直接使用Union明确指定类型 #

这是最直观的方式，直接列出所有允许的类型：

from typing import Union
import numpy as np

def func(l: Union[list, np.ndarray]) -> None:
    pass

优点：简单直接，类型检查工具能准确识别；缺点：如果后续需要支持更多类数组对象（如pandas Series），需要手动添加到Union中，扩展性一般。

2. 自定义Protocol协议 #

通过typing.Protocol定义一个包含你需要方法的协议，任何实现了这些方法的类型都会被自动兼容：

from typing import Protocol, TypeVar
import numpy as np

T = TypeVar('T')  # 用于指定元素类型

class ArrayLike(Protocol[T]):
    def __len__(self) -> int: ...
    def __getitem__(self, key: int) -> T: ...

def func(l: ArrayLike[int]) -> None:  # 指定元素为int类型，也可以省略用ArrayLike
    pass

这种方式灵活性最高，你可以根据实际需求添加更多方法（比如__iter__），同时支持所有符合协议的类数组对象，包括list、np.ndarray甚至自定义的类数组结构。主流类型检查工具（如mypy、pyright）都支持Protocol。

3. 使用numpy官方的ArrayLike类型 #

如果你主要处理数值类的类数组对象，numpy自带的numpy.typing.ArrayLike是更专业的选择，它已经涵盖了list、tuple、np.ndarray、pandas.Series等多种常用类数组类型：

import numpy as np
from numpy.typing import ArrayLike

def func(l: ArrayLike) -> None:
    pass

这个类型是numpy为数值计算场景专门设计的，能很好适配numpy的使用习惯，类型检查工具也能准确识别。

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