Type annotate tests (#5728)

* Type annotate lots of tests

* fixes for newer numpy version

* .

* .

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.