Change annotations to allow str keys (#5690)

* Change typing to allow str keys

* Change all incoming Mapping types

* Add in some annotated tests

* whatsnew

Author: max-sixty

doc/whats-new.rst

@@ -53,6 +53,8 @@ Internal Changes
   By `Deepak Cherian <https://github.com/dcherian>`_.
 - Explicit indexes refactor: decouple ``xarray.Index``` from ``xarray.Variable`` (:pull:`5636`).
   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>`_.
 - 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_pandas_roundtrip.py

@@ -28,7 +28,7 @@
 
 
 @st.composite
-def datasets_1d_vars(draw):
+def datasets_1d_vars(draw) -> xr.Dataset:
     """Generate datasets with only 1D variables
 
     Suitable for converting to pandas dataframes.
@@ -49,7 +49,7 @@ def datasets_1d_vars(draw):
 
 
 @given(st.data(), an_array)
-def test_roundtrip_dataarray(data, arr):
+def test_roundtrip_dataarray(data, arr) -> None:
     names = data.draw(
         st.lists(st.text(), min_size=arr.ndim, max_size=arr.ndim, unique=True).map(
             tuple
@@ -62,15 +62,15 @@ def test_roundtrip_dataarray(data, arr):
 
 
 @given(datasets_1d_vars())
-def test_roundtrip_dataset(dataset):
+def test_roundtrip_dataset(dataset) -> None:
     df = dataset.to_dataframe()
     assert isinstance(df, pd.DataFrame)
     roundtripped = xr.Dataset(df)
     xr.testing.assert_identical(dataset, roundtripped)
 
 
 @given(numeric_series, st.text())
-def test_roundtrip_pandas_series(ser, ix_name):
+def test_roundtrip_pandas_series(ser, ix_name) -> None:
     # Need to name the index, otherwise Xarray calls it 'dim_0'.
     ser.index.name = ix_name
     arr = xr.DataArray(ser)
@@ -87,7 +87,7 @@ def test_roundtrip_pandas_series(ser, ix_name):
 
 @pytest.mark.xfail
 @given(numeric_homogeneous_dataframe)
-def test_roundtrip_pandas_dataframe(df):
+def test_roundtrip_pandas_dataframe(df) -> None:
     # Need to name the indexes, otherwise Xarray names them 'dim_0', 'dim_1'.
     df.index.name = "rows"
     df.columns.name = "cols"

xarray/core/accessor_str.py

@@ -114,7 +114,7 @@ def _apply_str_ufunc(
     obj: Any,
     dtype: Union[str, np.dtype, Type] = None,
     output_core_dims: Union[list, tuple] = ((),),
-    output_sizes: Mapping[Hashable, int] = None,
+    output_sizes: Mapping[Any, int] = None,
     func_args: Tuple = (),
     func_kwargs: Mapping = {},
 ) -> Any:
@@ -227,7 +227,7 @@ def _apply(
         func: Callable,
         dtype: Union[str, np.dtype, Type] = None,
         output_core_dims: Union[list, tuple] = ((),),
-        output_sizes: Mapping[Hashable, int] = None,
+        output_sizes: Mapping[Any, int] = None,
         func_args: Tuple = (),
         func_kwargs: Mapping = {},
     ) -> Any:

xarray/core/common.py

@@ -818,7 +818,7 @@ def weighted(
 
     def rolling(
         self,
-        dim: Mapping[Hashable, int] = None,
+        dim: Mapping[Any, int] = None,
         min_periods: int = None,
         center: Union[bool, Mapping[Hashable, bool]] = False,
         **window_kwargs: int,
@@ -892,7 +892,7 @@ def rolling(
 
     def rolling_exp(
         self,
-        window: Mapping[Hashable, int] = None,
+        window: Mapping[Any, int] = None,
         window_type: str = "span",
         **window_kwargs,
     ):
@@ -933,7 +933,7 @@ def rolling_exp(
 
     def coarsen(
         self,
-        dim: Mapping[Hashable, int] = None,
+        dim: Mapping[Any, int] = None,
         boundary: str = "exact",
         side: Union[str, Mapping[Hashable, str]] = "left",
         coord_func: str = "mean",
@@ -1009,7 +1009,7 @@ def coarsen(
 
     def resample(
         self,
-        indexer: Mapping[Hashable, str] = None,
+        indexer: Mapping[Any, str] = None,
         skipna=None,
         closed: str = None,
         label: str = None,

xarray/core/computation.py

@@ -400,7 +400,7 @@ def apply_dict_of_variables_vfunc(
 
 
 def _fast_dataset(
-    variables: Dict[Hashable, Variable], coord_variables: Mapping[Hashable, Variable]
+    variables: Dict[Hashable, Variable], coord_variables: Mapping[Any, Variable]
 ) -> "Dataset":
     """Create a dataset as quickly as possible.
 

xarray/core/coordinates.py

@@ -158,7 +158,7 @@ def to_index(self, ordered_dims: Sequence[Hashable] = None) -> pd.Index:
 
         return pd.MultiIndex(level_list, code_list, names=names)
 
-    def update(self, other: Mapping[Hashable, Any]) -> None:
+    def update(self, other: Mapping[Any, Any]) -> None:
         other_vars = getattr(other, "variables", other)
         coords, indexes = merge_coords(
             [self.variables, other_vars], priority_arg=1, indexes=self.xindexes
@@ -270,7 +270,7 @@ def to_dataset(self) -> "Dataset":
         return self._data._copy_listed(names)
 
     def _update_coords(
-        self, coords: Dict[Hashable, Variable], indexes: Mapping[Hashable, Index]
+        self, coords: Dict[Hashable, Variable], indexes: Mapping[Any, Index]
     ) -> None:
         from .dataset import calculate_dimensions
 
@@ -333,7 +333,7 @@ def __getitem__(self, key: Hashable) -> "DataArray":
         return self._data._getitem_coord(key)
 
     def _update_coords(
-        self, coords: Dict[Hashable, Variable], indexes: Mapping[Hashable, Index]
+        self, coords: Dict[Hashable, Variable], indexes: Mapping[Any, Index]
     ) -> None:
         from .dataset import calculate_dimensions
 
@@ -376,7 +376,7 @@ def _ipython_key_completions_(self):
 
 
 def assert_coordinate_consistent(
-    obj: Union["DataArray", "Dataset"], coords: Mapping[Hashable, Variable]
+    obj: Union["DataArray", "Dataset"], coords: Mapping[Any, Variable]
 ) -> None:
     """Make sure the dimension coordinate of obj is consistent with coords.
 
@@ -394,7 +394,7 @@ def assert_coordinate_consistent(
 
 def remap_label_indexers(
     obj: Union["DataArray", "Dataset"],
-    indexers: Mapping[Hashable, Any] = None,
+    indexers: Mapping[Any, Any] = None,
     method: str = None,
     tolerance=None,
     **indexers_kwargs: Any,

xarray/core/dataarray.py

@@ -462,7 +462,7 @@ def _replace_maybe_drop_dims(
             )
         return self._replace(variable, coords, name, indexes=indexes)
 
-    def _overwrite_indexes(self, indexes: Mapping[Hashable, Any]) -> "DataArray":
+    def _overwrite_indexes(self, indexes: Mapping[Any, Any]) -> "DataArray":
         if not len(indexes):
             return self
         coords = self._coords.copy()
@@ -792,7 +792,7 @@ def attrs(self) -> Dict[Hashable, Any]:
         return self.variable.attrs
 
     @attrs.setter
-    def attrs(self, value: Mapping[Hashable, Any]) -> None:
+    def attrs(self, value: Mapping[Any, Any]) -> None:
         # Disable type checking to work around mypy bug - see mypy#4167
         self.variable.attrs = value  # type: ignore[assignment]
 
@@ -803,7 +803,7 @@ def encoding(self) -> Dict[Hashable, Any]:
         return self.variable.encoding
 
     @encoding.setter
-    def encoding(self, value: Mapping[Hashable, Any]) -> None:
+    def encoding(self, value: Mapping[Any, Any]) -> None:
         self.variable.encoding = value
 
     @property
@@ -1110,7 +1110,7 @@ def chunk(
 
     def isel(
         self,
-        indexers: Mapping[Hashable, Any] = None,
+        indexers: Mapping[Any, Any] = None,
         drop: bool = False,
         missing_dims: str = "raise",
         **indexers_kwargs: Any,
@@ -1193,7 +1193,7 @@ def isel(
 
     def sel(
         self,
-        indexers: Mapping[Hashable, Any] = None,
+        indexers: Mapping[Any, Any] = None,
         method: str = None,
         tolerance=None,
         drop: bool = False,
@@ -1498,7 +1498,7 @@ def reindex_like(
 
     def reindex(
         self,
-        indexers: Mapping[Hashable, Any] = None,
+        indexers: Mapping[Any, Any] = None,
         method: str = None,
         tolerance=None,
         copy: bool = True,
@@ -1591,7 +1591,7 @@ def reindex(
 
     def interp(
         self,
-        coords: Mapping[Hashable, Any] = None,
+        coords: Mapping[Any, Any] = None,
         method: str = "linear",
         assume_sorted: bool = False,
         kwargs: Mapping[str, Any] = None,
@@ -1815,7 +1815,7 @@ def rename(
             return self._replace(name=new_name_or_name_dict)
 
     def swap_dims(
-        self, dims_dict: Mapping[Hashable, Hashable] = None, **dims_kwargs
+        self, dims_dict: Mapping[Any, Hashable] = None, **dims_kwargs
     ) -> "DataArray":
         """Returns a new DataArray with swapped dimensions.
 
@@ -2333,7 +2333,7 @@ def drop(
 
     def drop_sel(
         self,
-        labels: Mapping[Hashable, Any] = None,
+        labels: Mapping[Any, Any] = None,
         *,
         errors: str = "raise",
         **labels_kwargs,
@@ -3163,7 +3163,7 @@ def diff(self, dim: Hashable, n: int = 1, label: Hashable = "upper") -> "DataArr
 
     def shift(
         self,
-        shifts: Mapping[Hashable, int] = None,
+        shifts: Mapping[Any, int] = None,
         fill_value: Any = dtypes.NA,
         **shifts_kwargs: int,
     ) -> "DataArray":
@@ -3210,7 +3210,7 @@ def shift(
 
     def roll(
         self,
-        shifts: Mapping[Hashable, int] = None,
+        shifts: Mapping[Any, int] = None,
         roll_coords: bool = None,
         **shifts_kwargs: int,
     ) -> "DataArray":
@@ -4433,7 +4433,7 @@ def argmax(
 
     def query(
         self,
-        queries: Mapping[Hashable, Any] = None,
+        queries: Mapping[Any, Any] = None,
         parser: str = "pandas",
         engine: str = None,
         missing_dims: str = "raise",