NumPy `__array_ufunc__` does not work with typing

What is your issue?

When using NumPy functions which have been patched to work with xarray objects using __array_ufunc__, typing tools like MyPy do not correctly calculate the return value.

For example, the function np.exp has been adapted by xarray to return a DataArray if it is passed a DataArray, so that code like the following will work:

import xarray as xr
import numpy as np

da = xr.DataArray([1, 2, 3, 4])
blah = np.exp(da).rename("blah")

This code creates an xr.DataArray, uses numpy to calculate the exp of its values and then, since it is returned as a xr.DataArray it can call xarray methods like rename on it.

However, running MyPy on this code gives the error:

error: "ndarray[Any, dtype[Any]]" has no attribute "rename"

This is because there is typing information from NumPy which claims that np.exp returns ndarray[Any, dtype[Any]].

Now, I'm unsure whether this is a bug in xarray not providing the typing information to the __array_ufunc__ code, a bug in NumPy for not having the flexibility to type these overrides or in MyPy for now allowing for overrides like this to happen but I wanted to star with the place where the observed error occurs.

Start by reproducing the MyPy error with the np.exp(DataArray).rename("blah") example, then inspect xarray's array_ufunc behavior alongside NumPy's np.exp typing. Done means the return type is inferred as an xarray DataArray and the rename call type-checks, with a regression test covering the example.