Type annotate tests (#5728)

* Type annotate lots of tests

* fixes for newer numpy version

* .

* .

Author: max-sixty

doc/whats-new.rst

@@ -63,6 +63,9 @@ Internal Changes
   By `Benoit Bovy <https://github.com/benbovy>`_.
 - Fix ``Mapping`` argument typing to allow mypy to pass on ``str`` keys (:pull:`5690`).
   By `Maximilian Roos <https://github.com/max-sixty>`_.
+- Annotate many of our tests, and fix some of the resulting typing errors. This will
+  also mean our typing annotations are tested as part of CI. (:pull:`5728`).
+  By `Maximilian Roos <https://github.com/max-sixty>`_.
 - Improve the performance of reprs for large datasets or dataarrays. (:pull:`5661`)
   By `Jimmy Westling <https://github.com/illviljan>`_.
 - Use isort's `float_to_top` config. (:pull:`5695`).

properties/test_encode_decode.py

@@ -25,7 +25,7 @@
 
 @pytest.mark.slow
 @given(st.data(), an_array)
-def test_CFMask_coder_roundtrip(data, arr):
+def test_CFMask_coder_roundtrip(data, arr) -> None:
     names = data.draw(
         st.lists(st.text(), min_size=arr.ndim, max_size=arr.ndim, unique=True).map(
             tuple
@@ -39,7 +39,7 @@ def test_CFMask_coder_roundtrip(data, arr):
 
 @pytest.mark.slow
 @given(st.data(), an_array)
-def test_CFScaleOffset_coder_roundtrip(data, arr):
+def test_CFScaleOffset_coder_roundtrip(data, arr) -> None:
     names = data.draw(
         st.lists(st.text(), min_size=arr.ndim, max_size=arr.ndim, unique=True).map(
             tuple

xarray/core/common.py

@@ -411,7 +411,7 @@ def get_index(self, key: Hashable) -> pd.Index:
             return pd.Index(range(self.sizes[key]), name=key)
 
     def _calc_assign_results(
-        self: C, kwargs: Mapping[Hashable, Union[T, Callable[[C], T]]]
+        self: C, kwargs: Mapping[Any, Union[T, Callable[[C], T]]]
     ) -> Dict[Hashable, T]:
         return {k: v(self) if callable(v) else v for k, v in kwargs.items()}
 
@@ -820,7 +820,7 @@ def rolling(
         self,
         dim: Mapping[Any, int] = None,
         min_periods: int = None,
-        center: Union[bool, Mapping[Hashable, bool]] = False,
+        center: Union[bool, Mapping[Any, bool]] = False,
         **window_kwargs: int,
     ):
         """
@@ -935,7 +935,7 @@ def coarsen(
         self,
         dim: Mapping[Any, int] = None,
         boundary: str = "exact",
-        side: Union[str, Mapping[Hashable, str]] = "left",
+        side: Union[str, Mapping[Any, str]] = "left",
         coord_func: str = "mean",
         **window_kwargs: int,
     ):
@@ -1520,7 +1520,7 @@ def __getitem__(self, value):
 def full_like(
     other: "Dataset",
     fill_value,
-    dtype: Union[DTypeLike, Mapping[Hashable, DTypeLike]] = None,
+    dtype: Union[DTypeLike, Mapping[Any, DTypeLike]] = None,
 ) -> "Dataset":
     ...
 

xarray/core/computation.py

@@ -367,7 +367,7 @@ def _as_variables_or_variable(arg):
 
 
 def _unpack_dict_tuples(
-    result_vars: Mapping[Hashable, Tuple[Variable, ...]], num_outputs: int
+    result_vars: Mapping[Any, Tuple[Variable, ...]], num_outputs: int
 ) -> Tuple[Dict[Hashable, Variable], ...]:
     out: Tuple[Dict[Hashable, Variable], ...] = tuple({} for _ in range(num_outputs))
     for name, values in result_vars.items():

xarray/core/coordinates.py

@@ -31,7 +31,7 @@
 _THIS_ARRAY = ReprObject("<this-array>")
 
 
-class Coordinates(Mapping[Hashable, "DataArray"]):
+class Coordinates(Mapping[Any, "DataArray"]):
     __slots__ = ()
 
     def __getitem__(self, key: Hashable) -> "DataArray":

xarray/core/dataarray.py

@@ -366,7 +366,7 @@ class DataArray(AbstractArray, DataWithCoords, DataArrayArithmetic):
     def __init__(
         self,
         data: Any = dtypes.NA,
-        coords: Union[Sequence[Tuple], Mapping[Hashable, Any], None] = None,
+        coords: Union[Sequence[Tuple], Mapping[Any, Any], None] = None,
         dims: Union[Hashable, Sequence[Hashable], None] = None,
         name: Hashable = None,
         attrs: Mapping = None,
@@ -788,7 +788,8 @@ def loc(self) -> _LocIndexer:
         return _LocIndexer(self)
 
     @property
-    def attrs(self) -> Dict[Hashable, Any]:
+    # Key type needs to be `Any` because of mypy#4167
+    def attrs(self) -> Dict[Any, Any]:
         """Dictionary storing arbitrary metadata with this array."""
         return self.variable.attrs
 
@@ -1068,7 +1069,7 @@ def chunk(
             int,
             Tuple[int, ...],
             Tuple[Tuple[int, ...], ...],
-            Mapping[Hashable, Union[None, int, Tuple[int, ...]]],
+            Mapping[Any, Union[None, int, Tuple[int, ...]]],
         ] = {},  # {} even though it's technically unsafe, is being used intentionally here (#4667)
         name_prefix: str = "xarray-",
         token: str = None,
@@ -1312,7 +1313,7 @@ def sel(
 
     def head(
         self,
-        indexers: Union[Mapping[Hashable, int], int] = None,
+        indexers: Union[Mapping[Any, int], int] = None,
         **indexers_kwargs: Any,
     ) -> "DataArray":
         """Return a new DataArray whose data is given by the the first `n`
@@ -1329,7 +1330,7 @@ def head(
 
     def tail(
         self,
-        indexers: Union[Mapping[Hashable, int], int] = None,
+        indexers: Union[Mapping[Any, int], int] = None,
         **indexers_kwargs: Any,
     ) -> "DataArray":
         """Return a new DataArray whose data is given by the the last `n`
@@ -1346,7 +1347,7 @@ def tail(
 
     def thin(
         self,
-        indexers: Union[Mapping[Hashable, int], int] = None,
+        indexers: Union[Mapping[Any, int], int] = None,
         **indexers_kwargs: Any,
     ) -> "DataArray":
         """Return a new DataArray whose data is given by each `n` value
@@ -1778,7 +1779,7 @@ def interp_like(
 
     def rename(
         self,
-        new_name_or_name_dict: Union[Hashable, Mapping[Hashable, Hashable]] = None,
+        new_name_or_name_dict: Union[Hashable, Mapping[Any, Hashable]] = None,
         **names: Hashable,
     ) -> "DataArray":
         """Returns a new DataArray with renamed coordinates or a new name.
@@ -1874,7 +1875,7 @@ def swap_dims(
 
     def expand_dims(
         self,
-        dim: Union[None, Hashable, Sequence[Hashable], Mapping[Hashable, Any]] = None,
+        dim: Union[None, Hashable, Sequence[Hashable], Mapping[Any, Any]] = None,
         axis=None,
         **dim_kwargs: Any,
     ) -> "DataArray":
@@ -1926,7 +1927,7 @@ def expand_dims(
 
     def set_index(
         self,
-        indexes: Mapping[Hashable, Union[Hashable, Sequence[Hashable]]] = None,
+        indexes: Mapping[Any, Union[Hashable, Sequence[Hashable]]] = None,
         append: bool = False,
         **indexes_kwargs: Union[Hashable, Sequence[Hashable]],
     ) -> "DataArray":
@@ -2014,7 +2015,7 @@ def reset_index(
 
     def reorder_levels(
         self,
-        dim_order: Mapping[Hashable, Sequence[int]] = None,
+        dim_order: Mapping[Any, Sequence[int]] = None,
         **dim_order_kwargs: Sequence[int],
     ) -> "DataArray":
         """Rearrange index levels using input order.
@@ -2049,7 +2050,7 @@ def reorder_levels(
 
     def stack(
         self,
-        dimensions: Mapping[Hashable, Sequence[Hashable]] = None,
+        dimensions: Mapping[Any, Sequence[Hashable]] = None,
         **dimensions_kwargs: Sequence[Hashable],
     ) -> "DataArray":
         """
@@ -3868,17 +3869,13 @@ def polyfit(
 
     def pad(
         self,
-        pad_width: Mapping[Hashable, Union[int, Tuple[int, int]]] = None,
+        pad_width: Mapping[Any, Union[int, Tuple[int, int]]] = None,
         mode: str = "constant",
-        stat_length: Union[
-            int, Tuple[int, int], Mapping[Hashable, Tuple[int, int]]
-        ] = None,
+        stat_length: Union[int, Tuple[int, int], Mapping[Any, Tuple[int, int]]] = None,
         constant_values: Union[
-            int, Tuple[int, int], Mapping[Hashable, Tuple[int, int]]
-        ] = None,
-        end_values: Union[
-            int, Tuple[int, int], Mapping[Hashable, Tuple[int, int]]
+            int, Tuple[int, int], Mapping[Any, Tuple[int, int]]
         ] = None,
+        end_values: Union[int, Tuple[int, int], Mapping[Any, Tuple[int, int]]] = None,
         reflect_type: str = None,
         **pad_width_kwargs: Any,
     ) -> "DataArray":

xarray/core/dataset.py

@@ -210,7 +210,7 @@ def calculate_dimensions(variables: Mapping[Any, Variable]) -> Dict[Hashable, in
 
 
 def merge_indexes(
-    indexes: Mapping[Hashable, Union[Hashable, Sequence[Hashable]]],
+    indexes: Mapping[Any, Union[Hashable, Sequence[Hashable]]],
     variables: Mapping[Any, Variable],
     coord_names: Set[Hashable],
     append: bool = False,
@@ -510,7 +510,7 @@ def _initialize_feasible(lb, ub):
     return param_defaults, bounds_defaults
 
 
-class DataVariables(Mapping[Hashable, "DataArray"]):
+class DataVariables(Mapping[Any, "DataArray"]):
     __slots__ = ("_dataset",)
 
     def __init__(self, dataset: "Dataset"):
@@ -2110,7 +2110,7 @@ def chunk(
         chunks: Union[
             int,
             str,
-            Mapping[Hashable, Union[None, int, str, Tuple[int, ...]]],
+            Mapping[Any, Union[None, int, str, Tuple[int, ...]]],
         ] = {},  # {} even though it's technically unsafe, is being used intentionally here (#4667)
         name_prefix: str = "xarray-",
         token: str = None,
@@ -2485,7 +2485,7 @@ def sel(
 
     def head(
         self,
-        indexers: Union[Mapping[Hashable, int], int] = None,
+        indexers: Union[Mapping[Any, int], int] = None,
         **indexers_kwargs: Any,
     ) -> "Dataset":
         """Returns a new dataset with the first `n` values of each array
@@ -2531,7 +2531,7 @@ def head(
 
     def tail(
         self,
-        indexers: Union[Mapping[Hashable, int], int] = None,
+        indexers: Union[Mapping[Any, int], int] = None,
         **indexers_kwargs: Any,
     ) -> "Dataset":
         """Returns a new dataset with the last `n` values of each array
@@ -2580,7 +2580,7 @@ def tail(
 
     def thin(
         self,
-        indexers: Union[Mapping[Hashable, int], int] = None,
+        indexers: Union[Mapping[Any, int], int] = None,
         **indexers_kwargs: Any,
     ) -> "Dataset":
         """Returns a new dataset with each array indexed along every `n`-th
@@ -3559,7 +3559,7 @@ def swap_dims(
 
     def expand_dims(
         self,
-        dim: Union[None, Hashable, Sequence[Hashable], Mapping[Hashable, Any]] = None,
+        dim: Union[None, Hashable, Sequence[Hashable], Mapping[Any, Any]] = None,
         axis: Union[None, int, Sequence[int]] = None,
         **dim_kwargs: Any,
     ) -> "Dataset":
@@ -3691,7 +3691,7 @@ def expand_dims(
 
     def set_index(
         self,
-        indexes: Mapping[Hashable, Union[Hashable, Sequence[Hashable]]] = None,
+        indexes: Mapping[Any, Union[Hashable, Sequence[Hashable]]] = None,
         append: bool = False,
         **indexes_kwargs: Union[Hashable, Sequence[Hashable]],
     ) -> "Dataset":
@@ -3789,7 +3789,7 @@ def reset_index(
 
     def reorder_levels(
         self,
-        dim_order: Mapping[Hashable, Sequence[int]] = None,
+        dim_order: Mapping[Any, Sequence[int]] = None,
         **dim_order_kwargs: Sequence[int],
     ) -> "Dataset":
         """Rearrange index levels using input order.
@@ -3858,7 +3858,7 @@ def _stack_once(self, dims, new_dim):
 
     def stack(
         self,
-        dimensions: Mapping[Hashable, Sequence[Hashable]] = None,
+        dimensions: Mapping[Any, Sequence[Hashable]] = None,
         **dimensions_kwargs: Sequence[Hashable],
     ) -> "Dataset":
         """
@@ -6929,17 +6929,13 @@ def polyfit(
 
     def pad(
         self,
-        pad_width: Mapping[Hashable, Union[int, Tuple[int, int]]] = None,
+        pad_width: Mapping[Any, Union[int, Tuple[int, int]]] = None,
         mode: str = "constant",
-        stat_length: Union[
-            int, Tuple[int, int], Mapping[Hashable, Tuple[int, int]]
-        ] = None,
+        stat_length: Union[int, Tuple[int, int], Mapping[Any, Tuple[int, int]]] = None,
         constant_values: Union[
-            int, Tuple[int, int], Mapping[Hashable, Tuple[int, int]]
-        ] = None,
-        end_values: Union[
-            int, Tuple[int, int], Mapping[Hashable, Tuple[int, int]]
+            int, Tuple[int, int], Mapping[Any, Tuple[int, int]]
         ] = None,
+        end_values: Union[int, Tuple[int, int], Mapping[Any, Tuple[int, int]]] = None,
         reflect_type: str = None,
         **pad_width_kwargs: Any,
     ) -> "Dataset":

xarray/core/indexes.py

@@ -34,7 +34,7 @@ class Index:
 
     @classmethod
     def from_variables(
-        cls, variables: Mapping[Hashable, "Variable"]
+        cls, variables: Mapping[Any, "Variable"]
     ) -> Tuple["Index", Optional[IndexVars]]:  # pragma: no cover
         raise NotImplementedError()
 
@@ -153,7 +153,7 @@ def __init__(self, array: Any, dim: Hashable):
         self.dim = dim
 
     @classmethod
-    def from_variables(cls, variables: Mapping[Hashable, "Variable"]):
+    def from_variables(cls, variables: Mapping[Any, "Variable"]):
         from .variable import IndexVariable
 
         if len(variables) != 1:
@@ -291,7 +291,7 @@ def _create_variables_from_multiindex(index, dim, level_meta=None):
 
 class PandasMultiIndex(PandasIndex):
     @classmethod
-    def from_variables(cls, variables: Mapping[Hashable, "Variable"]):
+    def from_variables(cls, variables: Mapping[Any, "Variable"]):
         if any([var.ndim != 1 for var in variables.values()]):
             raise ValueError("PandasMultiIndex only accepts 1-dimensional variables")
 
@@ -499,7 +499,7 @@ def isel_variable_and_index(
     name: Hashable,
     variable: "Variable",
     index: Index,
-    indexers: Mapping[Hashable, Union[int, slice, np.ndarray, "Variable"]],
+    indexers: Mapping[Any, Union[int, slice, np.ndarray, "Variable"]],
 ) -> Tuple["Variable", Optional[Index]]:
     """Index a Variable and an Index together.