将xarray DataArray传入numpy ufuncs时的类型标注问题

问题背景 #

我编写了一个带类型标注的函数：

import xarray as xr
import numpy as np

def compute_relative_azimuth(sat_azi: xr.DataArray, sun_azi: xr.DataArray) -> xr.DataArray:
    abs_diff = np.absolute(sun_azi - sat_azi)
    ssadiff = np.minimum(abs_diff, 360 - abs_diff)
    return ssadiff

该函数运行正常，但用mypy做类型检查时出现报错：

error: Incompatible return value type (got "ndarray[Any, dtype[Any]]", expected "DataArray")  [return-value]

补充说明：运行时np.minimum实际返回的是xr.DataArray，示例如下：

In [1]: import numpy as np

In [2]: import dask.array as da

In [3]: import xarray as xr

In [4]: a = xr.DataArray(da.zeros((5, 5)))

In [5]: b = xr.DataArray(da.zeros((5, 5)))

In [6]: np.minimum(a, b)
Out[6]:
<xarray.DataArray 'zeros_like-fe48a3d33a47c6e3255ffa508dd7dfba' (dim_0: 5,
                                                                 dim_1: 5)>
dask.array<minimum, shape=(5, 5), dtype=float64, chunksize=(5, 5), chunktype=numpy.ndarray>
Dimensions without coordinates: dim_0, dim_1

问题 #

除了忽略报错或强制类型转换外，有没有办法让mypy正确识别返回值类型？另外，xarray或numpy层面的长期解决方案是什么？

解决方案 #

即时解决方法 #

• 使用xarray内置方法：xarray.DataArray自带minimum实例方法，直接调用该方法替代np.minimum，mypy可以准确识别返回类型为xr.DataArray。修改后的代码：

import xarray as xr
import numpy as np

def compute_relative_azimuth(sat_azi: xr.DataArray, sun_azi: xr.DataArray) -> xr.DataArray:
    abs_diff = np.absolute(sun_azi - sat_azi)
    ssadiff = abs_diff.minimum(360 - abs_diff)
    return ssadiff

• 使用xarray封装的ufunc：xarray提供了xr.ufuncs模块，其中的函数专门适配xarray对象，用xr.ufuncs.minimum替换np.minimum，同样能让mypy正确推断返回类型。

长期解决方案 #

• xarray侧：xarray团队正在持续完善类型注解体系，后续版本会针对numpy ufunc的包装逻辑补充更精准的类型提示，让类型检查工具能自动识别numpy函数作用于DataArray时的返回类型。
• numpy侧：当前numpy的类型系统对自定义对象（如xarray.DataArray）的ufunc返回类型支持有限，未来可能会扩展类型注解机制，允许第三方库注册ufunc的类型映射关系，从而让mypy等工具可以正确推断返回类型。

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