More annotations in Dataset

doc/whats-new.rst

@@ -51,6 +51,15 @@ New functions/methods
   (:issue:`3026`).
   By `Julia Kent <https://github.com/jukent>`_.
 
+- 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
 ~~~~~~~~~~~~
 

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]
@@ -463,8 +480,8 @@ def broadcast(*args, exclude=None):
         a        (x, y) int64 1 1 2 2 3 3
         b        (x, y) int64 5 6 5 6 5 6
     """
-    from .dataarray import DataArray
-    from .dataset import Dataset
+    from .dataarray import DataArray  # noqa: F811
+    from .dataset import Dataset  # noqa: F811
 
     if exclude is None:
         exclude = set()

xarray/core/dataarray.py

@@ -1041,8 +1041,13 @@ def reindex_like(self, other: Union['DataArray', Dataset],
         align
         """
         indexers = reindex_like_indexers(self, other)
-        return self.reindex(method=method, tolerance=tolerance, copy=copy,
-                            fill_value=fill_value, **indexers)
+        return self.reindex(
+            indexers=indexers,
+            method=method,
+            tolerance=tolerance,
+            copy=copy,
+            fill_value=fill_value,
+        )
 
     def reindex(self, indexers: Optional[Mapping[Hashable, Any]] = None,
                 method: Optional[str] = None, tolerance=None,
@@ -1329,9 +1334,13 @@ def expand_dims(self, dim: Union[None, Hashable, Sequence[Hashable],
         ds = self._to_temp_dataset().expand_dims(dim, axis)
         return self._from_temp_dataset(ds)
 
-    def set_index(self, indexes: Optional[Mapping[Hashable, Any]] = None,
-                  append: bool = False, inplace: Optional[bool] = None,
-                  **indexes_kwargs: Any) -> Optional['DataArray']:
+    def set_index(
+        self,
+        indexes: Mapping[Hashable, Union[Hashable, Sequence[Hashable]]] = None,
+        append: bool = False,
+        inplace: bool = None,
+        **indexes_kwargs: Union[Hashable, Sequence[Hashable]]
+    ) -> Optional['DataArray']:
         """Set DataArray (multi-)indexes using one or more existing
         coordinates.
 
@@ -1370,9 +1379,12 @@ def set_index(self, indexes: Optional[Mapping[Hashable, Any]] = None,
         else:
             return self._replace(coords=coords)
 
-    def reset_index(self, dims_or_levels: Union[Hashable, Sequence[Hashable]],
-                    drop: bool = False, inplace: Optional[bool] = None
-                    ) -> Optional['DataArray']:
+    def reset_index(
+        self,
+        dims_or_levels: Union[Hashable, Sequence[Hashable]],
+        drop: bool = False,
+        inplace: bool = None
+    ) -> Optional['DataArray']:
         """Reset the specified index(es) or multi-index level(s).
 
         Parameters
@@ -1406,12 +1418,12 @@ def reset_index(self, dims_or_levels: Union[Hashable, Sequence[Hashable]],
         else:
             return self._replace(coords=coords)
 
-    def reorder_levels(self,
-                       dim_order: Optional[
-                           Mapping[Hashable, Sequence[int]]] = None,
-                       inplace: Optional[bool] = None,
-                       **dim_order_kwargs: Sequence[int]
-                       ) -> Optional['DataArray']:
+    def reorder_levels(
+        self,
+        dim_order: Mapping[Hashable, Sequence[int]] = None,
+        inplace: bool = None,
+        **dim_order_kwargs: Sequence[int]
+    ) -> Optional['DataArray']:
         """Rearrange index levels using input order.
 
         Parameters
@@ -1452,9 +1464,11 @@ def reorder_levels(self,
         else:
             return self._replace(coords=coords)
 
-    def stack(self, dimensions: Optional[
-            Mapping[Hashable, Sequence[Hashable]]] = None,
-            **dimensions_kwargs: Sequence[Hashable]) -> 'DataArray':
+    def stack(
+        self,
+        dimensions: Mapping[Hashable, Sequence[Hashable]] = None,
+        **dimensions_kwargs: Sequence[Hashable]
+    ) -> 'DataArray':
         """
         Stack any number of existing dimensions into a single new dimension.