带有 numpy ndarray 和 pandas 数据框的 mypy 重载函数（签名参数类型相同或更广泛）答案

【问题标题】：mypy overload function with numpy ndarray and pandas dataframe (signature parameter type(s) are the same or broader)带有 numpy ndarray 和 pandas 数据框的 mypy 重载函数（签名参数类型相同或更广泛）
【发布时间】：2020-11-22 14:47:15
【问题描述】：

我有一个函数可以对 numpy 数组或 pandas 数据帧执行一些算术运算（分位数归一化）。当你放入一个ndarray时，你应该得到一个ndarray，当你放入一个pandas数据框时，你应该得到一个数据框：

from typing import Union, overload


@overload
def quantile_normalize(data: pd.DataFrame) -> pd.DataFrame: ...
@overload
def quantile_normalize(data: np.ndarray) -> np.ndarray: ...


def quantile_normalize(data: Union[pd.DataFrame, np.ndarray]) -> Union[pd.DataFrame, np.ndarray]:
    pass

但是，当我尝试使用 mypy 进行测试时，它会抱怨：

qnorm/quantile_normalize.py:72: error: Overloaded function signature 2 will never be matched: signature 1's parameter type(s) are the same or broader

到目前为止，我发现与此问题相关的问题和答案似乎都与可选输入/输出和无类型有关。 pandas 数据框和 numpy 数组是相关的，但是它们应该可以通过 mypy 区分。

【问题讨论】：

您有numpy/pandas 的类型存根可用吗？例如。运行pip install data-science-types，删除因缺少numpy/pandas而忽略的类型提示并检查错误是否仍然存在。
@hoefling 确实解决了这个问题。请将此作为答案提交，我很乐意接受它:)

标签： python-3.x pandas numpy type-hinting mypy

【解决方案1】：

当库缺少类型提示时，每次导入都将解析为 Any。 numpy 和 pandas 都不符合 PEP 526（不提供任何类型提示）并且在 typeshed 中没有存根，因此 pd.DataFrame 和 np.ndarray 都将解析为 Any，因此两个重载都解析到def quantile_normalize(data: Any) -> Any: ...。要解决此问题，请为 numpy 和 pandas 添加存根。

要么使用现有的类型存根——我使用data-science-types（PyPI、GitHub）为numpy、pandas 和matplotlib 提供存根：

$ pip install data-science-types

现在pd.DataFrame 和np.ndarray 将在运行mypy 时正确解析。这还将在每个免费支持 PEP 526（例如 Visual Studio Code 或 WingIDE）的 IDE 中为您提供更好的代码完成。

或者，如果您不能/不想添加存根包，请编写您自己的最小存根，例如

# _typeshed/pandas/__init__.pyi

from typing import Any


def __getattr__(name: str) -> Any: ...  # incomplete

class DataFrame:
    def __getattr__(self, name: str) -> Any: ...  # incomplete

和

# _typeshed/numpy/__init__.pyi

from typing import Any


def __getattr__(name: str) -> Any: ...  # incomplete

class ndarray:
    def __getattr__(self, name: str) -> Any: ...  # incomplete

然后运行MYPYPATH=_typeshed mypy ...。

【讨论】：