Can reset_index return None?

[joooeey]

What happened:

The line

img = img.reset_index(["spatial"])

was not accepted by mypy. Mypy gave the following error:

redacted.py:69: error: Incompatible types in assignment (expression has type "Optional[DataArray]", variable has type "DataArray")

What you expected to happen:

I think that the return type of DataArray.reset_index should be DataArray, not Optional[DataArray]. Optional[Array] also includes None but the source code doesn't return None anywhere as far as I can see. Specifically, reset_index just redirects to _replace which returns a DataArray not Optional[DataArray] according to its type annotation.

Relevant code in xarray.core.dataarray.py:

class DataArray(AbstractArray, DataWithCoords, DataArrayArithmetic):
.
.
.
    def _replace(
        self,
        variable: Variable = None,
        coords=None,
        name: Union[Hashable, None, Default] = _default,
        indexes=None,
    ) -> "DataArray":
        if variable is None:
            variable = self.variable
        if coords is None:
            coords = self._coords
        if name is _default:
            name = self.name
        return type(self)(variable, coords, name=name, fastpath=True, indexes=indexes)
.
.
.
    def reset_index(
        self,
        dims_or_levels: Union[Hashable, Sequence[Hashable]],
        drop: bool = False,
    ) -> Optional["DataArray"]:
        """Reset the specified index(es) or multi-index level(s).
        Parameters
        ----------
        dims_or_levels : hashable or sequence of hashable
            Name(s) of the dimension(s) and/or multi-index level(s) that will
            be reset.
        drop : bool, optional
            If True, remove the specified indexes and/or multi-index levels
            instead of extracting them as new coordinates (default: False).
        Returns
        -------
        obj : DataArray
            Another dataarray, with this dataarray's data but replaced
            coordinates.
        See Also
        --------
        DataArray.set_index
        """
        coords, _ = split_indexes(
            dims_or_levels, self._coords, set(), self._level_coords, drop=drop
        )
        return self._replace(coords=coords)

Question:
Do I need to worry that there is any input for img, where img.reset_index would return None?

If I understand everything correctly and this is just a bug, I'm happy to put in a pull request to change the annotation.

Environment:

Output of xr.show_versions()

INSTALLED VERSIONS

commit: None
python: 3.9.4 (default, Apr 9 2021, 16:34:09)
[GCC 7.3.0]
python-bits: 64
OS: Linux
OS-release: 5.4.0-73-generic
machine: x86_64
processor: x86_64
byteorder: little
LC_ALL: None
LANG: en_US.UTF-8
LOCALE: ('en_US', 'UTF-8')
libhdf5: 1.12.0
libnetcdf: 4.7.4

xarray: 0.18.2
pandas: 1.2.4
numpy: 1.19.5
scipy: 1.6.1
netCDF4: 1.5.6
pydap: None
h5netcdf: None
h5py: 2.10.0
Nio: None
zarr: None
cftime: 1.5.0
nc_time_axis: None
PseudoNetCDF: None
rasterio: 1.2.4
cfgrib: None
iris: None
bottleneck: None
dask: installed
distributed: None
matplotlib: 3.4.2
cartopy: None
seaborn: None
numbagg: None
pint: None
setuptools: 52.0.0.post20210125
pip: 21.0.1
conda: None
pytest: 6.2.4
IPython: 7.24.1
sphinx: 4.0.2

[max-sixty]

Good spot @joooeey . I think you're right. Does changing it cause any mypy errors?

[joooeey]

Well, when I change the return type to "DataArray" in my local project the mypy error I had reported is gone. It also doesn't introduce any new errors.
I don't have the xarray dev environment installed on my machine. Next thing I would do would be to make the pull request via the Github web interface and see if the CI passes.

[max-sixty]

Great @joooeey

[keewis]

maybe Optional["DataArray"] is a remnant of the time we had the inplace parameter?