More annotations in Dataset (#3112)

* Let mypy discover xarray typing hints

* Docs tweak

* More annotations in Dataset

* Fix Python 3.5-specific check for OrderedDict

* typing in merge

* broadcast_like typing tweaks

* Changed redundant x: Optional[MyType] = None to x: MyType = None

* flake8

* What's New

doc/whats-new.rst

@@ -28,6 +28,15 @@ New functions/methods
   By `Deepak Cherian <https://github.com/dcherian>`_ and `David Mertz 
   <http://github.com/DavidMertz>`_.
 
+- The xarray package is now discoverably by mypy (although typing hints
+  coverage is not complete yet). mypy users can now remove from their setup.cfg
+  the lines::
+
+    [mypy-xarray]
+    ignore_missing_imports = True
+
+   By `Guido Imperiale <https://github.com/crusaderky>`_
+
 Enhancements
 ~~~~~~~~~~~~
 
@@ -70,7 +79,6 @@ New functions/methods
   (:issue:`3026`).
   By `Julia Kent <https://github.com/jukent>`_.
 
-
 Enhancements
 ~~~~~~~~~~~~
 

setup.py

@@ -104,4 +104,4 @@
       tests_require=TESTS_REQUIRE,
       url=URL,
       packages=find_packages(),
-      package_data={'xarray': ['tests/data/*']})
+      package_data={'xarray': ['py.typed', 'tests/data/*']})

xarray/core/alignment.py

@@ -3,7 +3,16 @@
 import warnings
 from collections import OrderedDict, defaultdict
 from contextlib import suppress
-from typing import Any, Mapping, Optional, Tuple
+from typing import (
+    Any,
+    Dict,
+    Hashable,
+    Mapping,
+    Optional,
+    Tuple,
+    Union,
+    TYPE_CHECKING,
+)
 
 import numpy as np
 import pandas as pd
@@ -13,6 +22,10 @@
 from .utils import is_dict_like, is_full_slice
 from .variable import IndexVariable, Variable
 
+if TYPE_CHECKING:
+    from .dataarray import DataArray
+    from .dataset import Dataset
+
 
 def _get_joiner(join):
     if join == 'outer':
@@ -169,8 +182,8 @@ def deep_align(objects, join='inner', copy=True, indexes=None,
 
     This function is not public API.
     """
-    from .dataarray import DataArray
-    from .dataset import Dataset
+    from .dataarray import DataArray  # noqa: F811
+    from .dataset import Dataset  # noqa: F811
 
     if indexes is None:
         indexes = {}
@@ -222,7 +235,10 @@ def is_alignable(obj):
     return out
 
 
-def reindex_like_indexers(target, other):
+def reindex_like_indexers(
+    target: Union['DataArray', 'Dataset'],
+    other: Union['DataArray', 'Dataset'],
+) -> Dict[Hashable, pd.Index]:
     """Extract indexers to align target with other.
 
     Not public API.
@@ -236,7 +252,8 @@ def reindex_like_indexers(target, other):
 
     Returns
     -------
-    Dict[Any, pandas.Index] providing indexes for reindex keyword arguments.
+    Dict[Hashable, pandas.Index] providing indexes for reindex keyword
+    arguments.
 
     Raises
     ------
@@ -310,7 +327,7 @@ def reindex_variables(
     new_indexes : OrderedDict
         Dict of indexes associated with the reindexed variables.
     """
-    from .dataarray import DataArray
+    from .dataarray import DataArray  # noqa: F811
 
     # create variables for the new dataset
     reindexed = OrderedDict()  # type: OrderedDict[Any, Variable]
@@ -407,8 +424,8 @@ def _get_broadcast_dims_map_common_coords(args, exclude):
 
 def _broadcast_helper(arg, exclude, dims_map, common_coords):
 
-    from .dataarray import DataArray
-    from .dataset import Dataset
+    from .dataarray import DataArray  # noqa: F811
+    from .dataset import Dataset  # noqa: F811
 
     def _set_dims(var):
         # Add excluded dims to a copy of dims_map

xarray/core/common.py

@@ -3,7 +3,7 @@
 from textwrap import dedent
 from typing import (
     Any, Callable, Hashable, Iterable, Iterator, List, Mapping, MutableMapping,
-    Optional, Tuple, TypeVar, Union)
+    Tuple, TypeVar, Union)
 
 import numpy as np
 import pandas as pd
@@ -101,7 +101,7 @@ def __int__(self: Any) -> int:
     def __complex__(self: Any) -> complex:
         return complex(self.values)
 
-    def __array__(self: Any, dtype: Optional[DTypeLike] = None) -> np.ndarray:
+    def __array__(self: Any, dtype: DTypeLike = None) -> np.ndarray:
         return np.asarray(self.values, dtype=dtype)
 
     def __repr__(self) -> str:
@@ -448,7 +448,7 @@ def pipe(self, func: Union[Callable[..., T], Tuple[Callable[..., T], str]],
             return func(self, *args, **kwargs)
 
     def groupby(self, group, squeeze: bool = True,
-                restore_coord_dims: Optional[bool] = None):
+                restore_coord_dims: bool = None):
         """Returns a GroupBy object for performing grouped operations.
 
         Parameters
@@ -501,7 +501,7 @@ def groupby(self, group, squeeze: bool = True,
     def groupby_bins(self, group, bins, right: bool = True, labels=None,
                      precision: int = 3, include_lowest: bool = False,
                      squeeze: bool = True,
-                     restore_coord_dims: Optional[bool] = None):
+                     restore_coord_dims: bool = None):
         """Returns a GroupBy object for performing grouped operations.
 
         Rather than using all unique values of `group`, the values are discretized
@@ -557,8 +557,8 @@ def groupby_bins(self, group, bins, right: bool = True, labels=None,
                                              'include_lowest':
                                                  include_lowest})
 
-    def rolling(self, dim: Optional[Mapping[Hashable, int]] = None,
-                min_periods: Optional[int] = None, center: bool = False,
+    def rolling(self, dim: Mapping[Hashable, int] = None,
+                min_periods: int = None, center: bool = False,
                 **window_kwargs: int):
         """
         Rolling window object.
@@ -621,7 +621,7 @@ def rolling(self, dim: Optional[Mapping[Hashable, int]] = None,
 
     def rolling_exp(
         self,
-        window: Optional[Mapping[Hashable, int]] = None,
+        window: Mapping[Hashable, int] = None,
         window_type: str = 'span',
         **window_kwargs
     ):
@@ -658,7 +658,7 @@ def rolling_exp(
 
         return self._rolling_exp_cls(self, window, window_type)
 
-    def coarsen(self, dim: Optional[Mapping[Hashable, int]] = None,
+    def coarsen(self, dim: Mapping[Hashable, int] = None,
                 boundary: str = 'exact',
                 side: Union[str, Mapping[Hashable, str]] = 'left',
                 coord_func: str = 'mean',
@@ -721,11 +721,11 @@ def coarsen(self, dim: Optional[Mapping[Hashable, int]] = None,
             self, dim, boundary=boundary, side=side,
             coord_func=coord_func)
 
-    def resample(self, indexer: Optional[Mapping[Hashable, str]] = None,
-                 skipna=None, closed: Optional[str] = None,
-                 label: Optional[str] = None,
-                 base: int = 0, keep_attrs: Optional[bool] = None,
-                 loffset=None, restore_coord_dims: Optional[bool] = None,
+    def resample(self, indexer: Mapping[Hashable, str] = None,
+                 skipna=None, closed: str = None,
+                 label: str = None,
+                 base: int = 0, keep_attrs: bool = None,
+                 loffset=None, restore_coord_dims: bool = None,
                  **indexer_kwargs: str):
         """Returns a Resample object for performing resampling operations.
 
@@ -1003,7 +1003,7 @@ def __getitem__(self, value):
         raise NotImplementedError
 
 
-def full_like(other, fill_value, dtype: Optional[DTypeLike] = None):
+def full_like(other, fill_value, dtype: DTypeLike = None):
     """Return a new object with the same shape and type as a given object.
 
     Parameters
@@ -1044,7 +1044,7 @@ def full_like(other, fill_value, dtype: Optional[DTypeLike] = None):
 
 
 def _full_like_variable(other, fill_value,
-                        dtype: Optional[DTypeLike] = None):
+                        dtype: DTypeLike = None):
     """Inner function of full_like, where other must be a variable
     """
     from .variable import Variable
@@ -1061,13 +1061,13 @@ def _full_like_variable(other, fill_value,
     return Variable(dims=other.dims, data=data, attrs=other.attrs)
 
 
-def zeros_like(other, dtype: Optional[DTypeLike] = None):
+def zeros_like(other, dtype: DTypeLike = None):
     """Shorthand for full_like(other, 0, dtype)
     """
     return full_like(other, 0, dtype)
 
 
-def ones_like(other, dtype: Optional[DTypeLike] = None):
+def ones_like(other, dtype: DTypeLike = None):
     """Shorthand for full_like(other, 1, dtype)
     """
     return full_like(other, 1, dtype)

xarray/core/computation.py

@@ -683,7 +683,7 @@ def apply_array_ufunc(func, *args, dask='forbidden'):
 def apply_ufunc(
     func: Callable,
     *args: Any,
-    input_core_dims: Optional[Sequence[Sequence]] = None,
+    input_core_dims: Sequence[Sequence] = None,
     output_core_dims: Optional[Sequence[Sequence]] = ((),),
     exclude_dims: AbstractSet = frozenset(),
     vectorize: bool = False,
@@ -693,8 +693,8 @@ def apply_ufunc(
     keep_attrs: bool = False,
     kwargs: Mapping = None,
     dask: str = 'forbidden',
-    output_dtypes: Optional[Sequence] = None,
-    output_sizes: Optional[Mapping[Any, int]] = None
+    output_dtypes: Sequence = None,
+    output_sizes: Mapping[Any, int] = None
 ) -> Any:
     """Apply a vectorized function for unlabeled arrays on xarray objects.