为何mypy处理numpy标量/数组运算时返回Any类型报错？

问题描述 #

我疑惑为何mypy在处理numpy相关代码时，频繁抛出“Returning Any from function declared to return ...”错误。如下示例中，a是数组，x是标量，np.exp(a*x)显然只会返回数组，即便x为np.inf结果仍是数组，为何mypy认为该函数可能返回Any类型？

import numpy as np
from numpy.typing import NDArray

def func(a: NDArray[np.float64], x: float) -> NDArray[np.float64]:
    return np.exp(a*x)

报错信息 #

使用mypy --warn-return-any [文件名]检查上述代码，会出现如下报错：

a.py:5: error: Returning Any from function declared to return "ndarray[Any, dtype[floating[_64Bit]]]" [no-any-return]

环境版本 #

• python 3.10.8
• mypy 0.991
• numpy 1.24.0

原因分析 #

这个问题源于numpy旧版本类型注解与mypy的兼容性不足：

1. numpy类型存根的局限性：numpy 1.24.0中np.exp的类型定义没有完全覆盖所有输入组合的精确返回类型。当输入是NDArray[float64]与Python原生float的乘积时，类型存根无法明确标注返回类型，导致mypy只能推断为Any。
2. 标量与数组运算的类型盲区：Python原生float和numpy数组的运算结果，在旧版numpy的类型定义中未被精确描述，mypy无法确定a*x的具体类型，进而无法识别np.exp的返回类型。

解决方法 #

• 升级numpy版本：numpy 1.25及以上版本大幅优化了类型注解，修复了大量类型推断问题，升级后mypy可正确识别返回类型。
• 显式标注中间结果：如果暂时无法升级，给中间运算结果添加类型注解，帮助mypy完成推断：

  import numpy as np
  from numpy.typing import NDArray
  
  def func(a: NDArray[np.float64], x: float) -> NDArray[np.float64]:
      product: NDArray[np.float64] = a * x
      return np.exp(product)
  

• 替换为numpy标量类型：将x的类型从Python原生float改为np.float64，让运算和函数调用的类型推断更准确：

  import numpy as np
  from numpy.typing import NDArray
  
  def func(a: NDArray[np.float64], x: np.float64) -> NDArray[np.float64]:
      return np.exp(a*x)
  

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