Trying to resolve mypy issues

Author: jsignell

xarray/core/alignment.py

@@ -113,7 +113,7 @@ class Aligner(Generic[T_Alignable]):
     objects: tuple[T_Alignable, ...]
     results: tuple[T_Alignable, ...]
     objects_matching_indexes: tuple[dict[MatchingIndexKey, Index], ...]
-    join: str
+    join: str | CombineKwargDefault
     exclude_dims: frozenset[Hashable]
     exclude_vars: frozenset[Hashable]
     copy: bool
@@ -133,7 +133,7 @@ class Aligner(Generic[T_Alignable]):
     def __init__(
         self,
         objects: Iterable[T_Alignable],
-        join: str = "inner",
+        join: str | CombineKwargDefault = "inner",
         indexes: Mapping[Any, Any] | None = None,
         exclude_dims: str | Iterable[Hashable] = frozenset(),
         exclude_vars: Iterable[Hashable] = frozenset(),
@@ -146,7 +146,14 @@ def __init__(
         self.objects = tuple(objects)
         self.objects_matching_indexes = ()
 
-        if join not in ["inner", "outer", "override", "exact", "left", "right"]:
+        if not isinstance(join, CombineKwargDefault) and join not in [
+            "inner",
+            "outer",
+            "override",
+            "exact",
+            "left",
+            "right",
+        ]:
             raise ValueError(f"invalid value for join: {join}")
         self.join = join
 
@@ -618,7 +625,7 @@ def align(
     obj1: T_Obj1,
     /,
     *,
-    join: JoinOptions = "inner",
+    join: JoinOptions | CombineKwargDefault = "inner",
     copy: bool = True,
     indexes=None,
     exclude: str | Iterable[Hashable] = frozenset(),
@@ -632,7 +639,7 @@ def align(
     obj2: T_Obj2,
     /,
     *,
-    join: JoinOptions = "inner",
+    join: JoinOptions | CombineKwargDefault = "inner",
     copy: bool = True,
     indexes=None,
     exclude: str | Iterable[Hashable] = frozenset(),
@@ -647,7 +654,7 @@ def align(
     obj3: T_Obj3,
     /,
     *,
-    join: JoinOptions = "inner",
+    join: JoinOptions | CombineKwargDefault = "inner",
     copy: bool = True,
     indexes=None,
     exclude: str | Iterable[Hashable] = frozenset(),
@@ -663,7 +670,7 @@ def align(
     obj4: T_Obj4,
     /,
     *,
-    join: JoinOptions = "inner",
+    join: JoinOptions | CombineKwargDefault = "inner",
     copy: bool = True,
     indexes=None,
     exclude: str | Iterable[Hashable] = frozenset(),
@@ -680,7 +687,7 @@ def align(
     obj5: T_Obj5,
     /,
     *,
-    join: JoinOptions = "inner",
+    join: JoinOptions | CombineKwargDefault = "inner",
     copy: bool = True,
     indexes=None,
     exclude: str | Iterable[Hashable] = frozenset(),
@@ -691,7 +698,7 @@ def align(
 @overload
 def align(
     *objects: T_Alignable,
-    join: JoinOptions = "inner",
+    join: JoinOptions | CombineKwargDefault = "inner",
     copy: bool = True,
     indexes=None,
     exclude: str | Iterable[Hashable] = frozenset(),
@@ -701,7 +708,7 @@ def align(
 
 def align(
     *objects: T_Alignable,
-    join: JoinOptions = "inner",
+    join: JoinOptions | CombineKwargDefault = "inner",
     copy: bool = True,
     indexes=None,
     exclude: str | Iterable[Hashable] = frozenset(),

xarray/core/combine.py

@@ -269,7 +269,7 @@ def _combine_all_along_first_dim(
     dim,
     data_vars,
     coords,
-    compat: CompatOptions,
+    compat: CompatOptions | CombineKwargDefault,
     fill_value,
     join: JoinOptions | CombineKwargDefault,
     combine_attrs: CombineAttrsOptions,
@@ -298,7 +298,7 @@ def _combine_all_along_first_dim(
 def _combine_1d(
     datasets,
     concat_dim,
-    compat: CompatOptions,
+    compat: CompatOptions | CombineKwargDefault,
     data_vars,
     coords,
     fill_value,
@@ -365,7 +365,7 @@ def _nested_combine(
         return Dataset()
 
     if isinstance(concat_dim, str | DataArray) or concat_dim is None:
-        concat_dim = [concat_dim]  # type: ignore[assignment]
+        concat_dim = [concat_dim]
 
     # Arrange datasets for concatenation
     # Use information from the shape of the user input

xarray/core/concat.py

@@ -65,7 +65,7 @@ def concat(
     compat: CompatOptions | CombineKwargDefault = _COMPAT_CONCAT_DEFAULT,
     positions: Iterable[Iterable[int]] | None = None,
     fill_value: object = dtypes.NA,
-    join: JoinOptions | None = None,
+    join: JoinOptions | CombineKwargDefault = _JOIN_DEFAULT,
     combine_attrs: CombineAttrsOptions = "override",
     create_index_for_new_dim: bool = True,
 ) -> T_DataArray: ...
@@ -334,7 +334,7 @@ def _calc_concat_over(
     datasets,
     dim,
     dim_names,
-    data_vars: T_DataVars,
+    data_vars: T_DataVars | CombineKwargDefault,
     coords,
     compat,
 ):
@@ -485,7 +485,7 @@ def process_subset_opt(opt, subset):
                     )
             concat_over.update(opt)
 
-    warnings = []
+    warnings: list[str] = []
     process_subset_opt(data_vars, "data_vars")
     process_subset_opt(coords, "coords")
 
@@ -534,7 +534,7 @@ def _dataset_concat(
     datasets: Iterable[T_Dataset],
     dim: str | T_Variable | T_DataArray | pd.Index,
     data_vars: T_DataVars | CombineKwargDefault,
-    coords: str | list[str] | CombineKwargDefault,
+    coords: str | list[Hashable] | CombineKwargDefault,
     compat: CompatOptions | CombineKwargDefault,
     positions: Iterable[Iterable[int]] | None,
     fill_value: Any,
@@ -780,7 +780,7 @@ def _dataarray_concat(
     arrays: Iterable[T_DataArray],
     dim: str | T_Variable | T_DataArray | pd.Index,
     data_vars: T_DataVars | CombineKwargDefault,
-    coords: str | list[str] | CombineKwargDefault,
+    coords: str | list[Hashable] | CombineKwargDefault,
     compat: CompatOptions | CombineKwargDefault,
     positions: Iterable[Iterable[int]] | None,
     fill_value: object,

xarray/core/merge.py

@@ -527,7 +527,7 @@ def coerce_pandas_values(objects: Iterable[CoercibleMapping]) -> list[DatasetLik
 def _get_priority_vars_and_indexes(
     objects: Sequence[DatasetLike],
     priority_arg: int | None,
-    compat: CompatOptions = "equals",
+    compat: CompatOptions | CombineKwargDefault = "equals",
 ) -> dict[Hashable, MergeElement]:
     """Extract the priority variable from a list of mappings.
 

xarray/tests/test_backends.py

@@ -14,7 +14,7 @@
 import uuid
 import warnings
 from collections.abc import Generator, Iterator, Mapping
-from contextlib import ExitStack, nullcontext
+from contextlib import ExitStack
 from io import BytesIO
 from os import listdir
 from pathlib import Path
@@ -4614,12 +4614,16 @@ def test_open_mfdataset_dataset_combine_attrs(
             with set_options(
                 use_new_combine_kwarg_defaults=use_new_combine_kwarg_defaults
             ):
-                warning = (
+                warning: contextlib.AbstractContextManager = (
                     pytest.warns(FutureWarning)
                     if not use_new_combine_kwarg_defaults
-                    else nullcontext()
+                    else contextlib.nullcontext()
+                )
+                error: contextlib.AbstractContextManager = (
+                    pytest.raises(xr.MergeError)
+                    if expect_error
+                    else contextlib.nullcontext()
                 )
-                error = pytest.raises(xr.MergeError) if expect_error else nullcontext()
                 with warning:
                     with error:
                         with xr.open_mfdataset(
@@ -4785,13 +4789,15 @@ def test_open_mfdataset_warns_when_kwargs_set_to_different(
                 xr.concat([ds1, ds2], dim="t", **kwargs)
 
             with set_options(use_new_combine_kwarg_defaults=False):
-                if "data_vars" not in kwargs:
-                    expectation = pytest.warns(
+                expectation: contextlib.AbstractContextManager = (
+                    pytest.warns(
                         FutureWarning,
                         match="will change from data_vars='all'",
                     )
-                else:
-                    expectation = nullcontext()
+                    if "data_vars" not in kwargs
+                    else contextlib.nullcontext()
+                )
+
                 with pytest.warns(
                     FutureWarning,
                     match="will change from compat='equals'",

xarray/tests/test_concat.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Callable
-from contextlib import nullcontext
+from contextlib import AbstractContextManager, nullcontext
 from copy import deepcopy
 from typing import TYPE_CHECKING, Any, Literal
 
@@ -1476,13 +1476,14 @@ def test_concat_coords_kwarg(
                 FutureWarning,
                 match="will change from data_vars='all' to data_vars='minimal'",
             ):
-                if coords == "different":
-                    expectation = pytest.warns(
+                expectation: AbstractContextManager = (
+                    pytest.warns(
                         FutureWarning,
                         match="will change from compat='equals' to compat='override'",
                     )
-                else:
-                    expectation = nullcontext()
+                    if coords == "different"
+                    else nullcontext()
+                )
                 with expectation:
                     old = concat(datasets, data["dim1"], coords=coords)
 

xarray/util/deprecation_helpers.py

@@ -148,8 +148,11 @@ def wrapper(*args, **kwargs):
     return wrapper  # type: ignore[return-value]
 
 
-class CombineKwargDefault(ReprObject):
-    """Object that handles deprecation cycle for kwarg default values."""
+class CombineKwargDefault:
+    """Object that handles deprecation cycle for kwarg default values.
+
+    Similar to ReprObject
+    """
 
     _old: str
     _new: str
@@ -160,22 +163,29 @@ def __init__(self, *, name: str, old: str, new: str):
         self._old = old
         self._new = new
 
+    def __repr__(self) -> str:
+        return self._value
+
     def __eq__(self, other: ReprObject | Any) -> bool:
-        # TODO: What type can other be? ArrayLike?
         return (
             self._value == other._value
             if isinstance(other, ReprObject)
             else self._value == other
         )
 
     @property
-    def _value(self):
+    def _value(self) -> str:
         return self._new if OPTIONS["use_new_combine_kwarg_defaults"] else self._old
 
     def __hash__(self) -> int:
         return hash(self._value)
 
-    def warning_message(self, message: str, recommend_set_options: bool = True):
+    def __dask_tokenize__(self) -> object:
+        from dask.base import normalize_token
+
+        return normalize_token((type(self), self._value))
+
+    def warning_message(self, message: str, recommend_set_options: bool = True) -> str:
         if recommend_set_options:
             recommendation = (
                 " To opt in to new defaults and get rid of these warnings now "
@@ -194,7 +204,7 @@ def warning_message(self, message: str, recommend_set_options: bool = True):
             + recommendation
         )
 
-    def error_message(self):
+    def error_message(self) -> str:
         return (
             f" Error might be related to new default ({self._name}={self._new!r}). "
             f"Previously the default was {self._name}={self._old!r}. "