Improve typehints of xr.Dataset.__getitem__

[nbren12]

To resolve some common type-related errors, this PR adds some overload type hints to Dataset.__getitem__. Now mypy can correctly infer that hashable inputs return DataArrays.

• Closes Improving typing of xr.Dataset.__getitem__ #4125
• Passes isort -rc . && black . && mypy . && flake8

[pep8speaks]

Hello @nbren12! Thanks for updating this PR. We checked the lines you've touched for PEP 8 issues, and found:

There are currently no PEP 8 issues detected in this Pull Request. Cheers! 🍻

Comment last updated at 2020-06-13 17:43:56 UTC

[mathause]

The mypy check throws an error: xarray/core/dataset.py:1250: error: Overloaded function signature 2 will never be matched: signature 1's parameter type(s) are the same I think you can ignore the other failures. Should that be:

    @overload
    def __getitem__(self, key: Hashable) -> DataArray:
        ...

    @overload
    def __getitem__(self, key: Iterable[Hashable]) -> "Dataset":
        ...

? Only guessing, though - it was Any so there may be more options. @crusaderky

[nbren12]

@mathause On further consideration, I think it might not be possible to get this to work. This method has three behaviors:

• Mapping -> Dataset
• Hashable -> DataArray
• else (List): -> Dataset

With my limited understanding of mypy, I think that any two of these is supported by overload, but I'm not sure it's possible to support all 3. I tried several different options, but maybe I am missing something.

Would a good middle ground be something like this?

• Hashable -> DataArray
• Any -> Union[DataArray, Dataset]

I think this would work since both the input/outputs of the first one are subtypes of the second one. It's not a complete solution, but it would solve the most common problem of ds['a'] returning a union type rather than a DataArray.

[nbren12]

Okay. Assuming the tests pass, I think this is ready for review. I tried adding a test, but mypy didn't seem to find problems even with code that I know doesn't work (e.g. 'a'+ 1). Is there some strategy for testing tricky type hints like this?

In any case, this code does work:

$ cat test_mypy.py                                                                                                                                                                                                    (fv3net) 
import xarray as xr
ds = xr.Dataset({"a": ()})


arr = ds['a']
union_obj = ds[['a']]

reveal_locals()
$ mypy test_mypy.py                                                                                                                                                                                                   (fv3net) 
test_mypy.py:8: note: Revealed local types are:
test_mypy.py:8: note:     arr: xarray.core.dataarray.DataArray
test_mypy.py:8: note:     ds: xarray.core.dataset.Dataset
test_mypy.py:8: note:     union_obj: Union[xarray.core.dataarray.DataArray, xarray.core.dataset.Dataset]

[mathause]

Seems this was already discussed in GH3210 (comment) and see also the TODO:

xarray/xarray/core/dataset.py

Lines 1244 to 1250 in 8f688ea

	def __getitem__(self, key: Any) -> "Union[DataArray, Dataset]":
	"""Access variables or coordinates this dataset as a
	:py:class:`~xarray.DataArray`.
	Indexing with a list of names will return a new ``Dataset`` object.
	"""
	# TODO(shoyer): type this properly: https://github.com/python/mypy/issues/7328

(although it is not entirely clear to me whether this is actually fixed or not)

[max-sixty]

Nice find @mathause ; I remember that discussion now.

As @nbren12 said, this doesn't go all the way given mypy's restrictions, but seems like a dominant improvement.

Test failures are unrelated.

Thanks @nbren12 !

[nbren12]

No problem! I think I am done with this one unless you think its important that I document or test this somehow. Can someone review it?

[crusaderky]

I took the liberty to rework it, please have a look
Test script:

from typing import Hashable, Mapping
import xarray
ds: xarray.Dataset


class D(Hashable, Mapping):
    def __hash__(self): ...
    def __getitem__(self, item): ...
    def __iter__(self): ...
    def __len__(self): ...

reveal_type(ds["foo"])
reveal_type(ds[["foo", "bar"]])
reveal_type(ds[{}])
reveal_type(ds[D()])

mypy output:

t1.py:12: note: Revealed type is 'xarray.core.dataarray.DataArray'
t1.py:13: note: Revealed type is 'xarray.core.dataset.Dataset'
t1.py:14: note: Revealed type is 'xarray.core.dataset.Dataset'
t1.py:15: note: Revealed type is 'xarray.core.dataset.Dataset'

[shoyer]

It looks like this is fixed now and can be removed? Or perhaps we more it below above the third @overload and add a comment that Any means list?

[crusaderky]

I don't think it's fixed. Specifically, mypy can't deal in the signature with an overload of Mapping and Hashable. Curiously however, once you add #type: ignore to the overloaded signature, the actual type inspection works just fine (see my test script above).

[nbren12]

@crusaderky Thanks for the re-work. For my own benefit, could you explain why that code worked? I remember writing something very similar, and running into mypy errors. My understanding of how mypy intreprets overload seems incomplete.

[crusaderky]

@nbren12 it seems to me that mypy is being overly aggressive when parsing the hinted code (hence why I had to put # type: ignore on it) but it is being more lax when the same code is invoked somewhere else like in my test script. Overall I suspect it may be fragile and break in future mypy versions...