WIP: typing for DataArray/Dataset (#2929)

* typing for DataArray/Dataset (incomplete)

* Finished dataarray.py

* Bugfixes

* Bugfixes

* trivial

* Dataset (partial)

* Bugfixes

* Minor tweaks

* final changes to dataset & dataarray for mypy

* mypy str conflict

* remove temp ptw setting

* OrderedDict-related fixes

* fix tests re dict

* fix mypy re dict

* merge issue re encoding kwarg

* weird python35 syntax error

* Hashable -> Any for Dicts

* DataArray names as Hashable

* resolve some __default vs None issues

* whats-new

* just specify Tuple at the moment

* remove unneeded cast, make arg formatting more black-like

* adjust init args based on @shoyer review

* add Tuple[Tuple] to dataset init types

* tuples not all the way down

Author: crusaderky

.gitignore

@@ -35,6 +35,7 @@ pip-log.txt
 .tox
 nosetests.xml
 .cache
+.dmypy.json
 .mypy_cache
 .ropeproject/
 .tags*

doc/whats-new.rst

@@ -88,6 +88,10 @@ Enhancements
   that allows ignoring errors if a passed label or dimension is not in the dataset
   (:issue:`2994`).
   By `Andrew Ross <https://github.com/andrew-c-ross>`_.
+- Argument and return types are added to most methods on ``DataArray`` and ``Dataset``,
+  allowing static type checking both within xarray and external libraries.
+  Type checking with ``mypy`` is enabled in CI (though not required yet).
+  By `Guido Imperiale <https://github.com/crusaderky>`_ and `Maximilian Roos <https://github.com/max-sixty>`_. 
 
 
 Bug fixes

setup.cfg

@@ -111,4 +111,4 @@ tag_prefix = v
 parentdir_prefix = xarray-
 
 [aliases]
-test = pytest
+test = pytest

xarray/core/common.py

@@ -10,6 +10,7 @@
 
 from . import dtypes, duck_array_ops, formatting, ops
 from .arithmetic import SupportsArithmetic
+from .npcompat import DTypeLike
 from .options import _get_keep_attrs
 from .pycompat import dask_array_type
 from .rolling_exp import RollingExp
@@ -100,8 +101,7 @@ def __int__(self: Any) -> int:
     def __complex__(self: Any) -> complex:
         return complex(self.values)
 
-    def __array__(self: Any, dtype: Union[str, np.dtype, None] = None
-                  ) -> np.ndarray:
+    def __array__(self: Any, dtype: Optional[DTypeLike] = None) -> np.ndarray:
         return np.asarray(self.values, dtype=dtype)
 
     def __repr__(self) -> str:
@@ -997,7 +997,7 @@ def __exit__(self, exc_type, exc_value, traceback) -> None:
         self.close()
 
 
-def full_like(other, fill_value, dtype: Union[str, np.dtype, None] = None):
+def full_like(other, fill_value, dtype: Optional[DTypeLike] = None):
     """Return a new object with the same shape and type as a given object.
 
     Parameters
@@ -1038,7 +1038,7 @@ def full_like(other, fill_value, dtype: Union[str, np.dtype, None] = None):
 
 
 def _full_like_variable(other, fill_value,
-                        dtype: Union[str, np.dtype, None] = None):
+                        dtype: Optional[DTypeLike] = None):
     """Inner function of full_like, where other must be a variable
     """
     from .variable import Variable
@@ -1055,19 +1055,19 @@ def _full_like_variable(other, fill_value,
     return Variable(dims=other.dims, data=data, attrs=other.attrs)
 
 
-def zeros_like(other, dtype: Union[str, np.dtype, None] = None):
+def zeros_like(other, dtype: Optional[DTypeLike] = None):
     """Shorthand for full_like(other, 0, dtype)
     """
     return full_like(other, 0, dtype)
 
 
-def ones_like(other, dtype: Union[str, np.dtype, None] = None):
+def ones_like(other, dtype: Optional[DTypeLike] = None):
     """Shorthand for full_like(other, 1, dtype)
     """
     return full_like(other, 1, dtype)
 
 
-def is_np_datetime_like(dtype: Union[str, np.dtype]) -> bool:
+def is_np_datetime_like(dtype: DTypeLike) -> bool:
     """Check if a dtype is a subclass of the numpy datetime types
     """
     return (np.issubdtype(dtype, np.datetime64) or

xarray/core/computation.py

@@ -4,7 +4,6 @@
 import functools
 import itertools
 import operator
-import typing
 from collections import Counter, OrderedDict
 from distutils.version import LooseVersion
 from typing import (

xarray/core/coordinates.py

@@ -1,6 +1,7 @@
 import collections.abc
 from collections import OrderedDict
 from contextlib import contextmanager
+from typing import Hashable, Iterable, Iterator, Union
 
 import pandas as pd
 
@@ -9,6 +10,11 @@
     expand_and_merge_variables, merge_coords, merge_coords_for_inplace_math)
 from .utils import Frozen, ReprObject, either_dict_or_kwargs
 from .variable import Variable
+from .pycompat import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .dataarray import DataArray
+    from .dataset import Dataset
 
 # Used as the key corresponding to a DataArray's variable when converting
 # arbitrary DataArray objects to datasets
@@ -258,21 +264,20 @@ def _ipython_key_completions_(self):
         return self._data._ipython_key_completions_()
 
 
-class LevelCoordinatesSource:
+class LevelCoordinatesSource(Iterable[Hashable]):
     """Iterator for MultiIndex level coordinates.
 
     Used for attribute style lookup with AttrAccessMixin. Not returned directly
     by any public methods.
     """
-
-    def __init__(self, data_object):
+    def __init__(self, data_object: 'Union[DataArray, Dataset]'):
         self._data = data_object
 
     def __getitem__(self, key):
         # not necessary -- everything here can already be found in coords.
-        raise KeyError
+        raise KeyError()
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[Hashable]:
         return iter(self._data._level_coords)