Use typevars rather than specific unions

Also from the chars of pydata#8208 -- this uses the TypeVars we define, which hopefully sets a standard and removes ambiguity for how to type functions. It also allows subclasses.

Where possible, it uses `T_Xarray`, otherwise `T_DataArrayOrSet` where it's not possible for `mypy` to narrow the concrete type down.

xarray/core/alignment.py

@@ -19,13 +19,11 @@
     indexes_all_equal,
     safe_cast_to_index,
 )
-from xarray.core.types import T_Alignable
+from xarray.core.types import T_Alignable, T_Xarray
 from xarray.core.utils import is_dict_like, is_full_slice
 from xarray.core.variable import Variable, as_compatible_data, calculate_dimensions
 
 if TYPE_CHECKING:
-    from xarray.core.dataarray import DataArray
-    from xarray.core.dataset import Dataset
     from xarray.core.types import JoinOptions, T_DataArray, T_Dataset, T_DuckArray
 
 
@@ -903,7 +901,7 @@ def reindex(
 
 def reindex_like(
     obj: T_Alignable,
-    other: Dataset | DataArray,
+    other: T_Xarray,
     method: str | None = None,
     tolerance: int | float | Iterable[int | float] | None = None,
     copy: bool = True,

xarray/core/coordinates.py

@@ -23,7 +23,7 @@
     create_default_index_implicit,
 )
 from xarray.core.merge import merge_coordinates_without_align, merge_coords
-from xarray.core.types import Self, T_DataArray
+from xarray.core.types import Self, T_DataArray, T_Xarray
 from xarray.core.utils import (
     Frozen,
     ReprObject,
@@ -915,9 +915,7 @@ def drop_indexed_coords(
     return Coordinates._construct_direct(coords=new_variables, indexes=new_indexes)
 
 
-def assert_coordinate_consistent(
-    obj: T_DataArray | Dataset, coords: Mapping[Any, Variable]
-) -> None:
+def assert_coordinate_consistent(obj: T_Xarray, coords: Mapping[Any, Variable]) -> None:
     """Make sure the dimension coordinate of obj is consistent with coords.
 
     obj: DataArray or Dataset

xarray/core/dataarray.py

@@ -41,6 +41,7 @@
 from xarray.core.indexing import is_fancy_indexer, map_index_queries
 from xarray.core.merge import PANDAS_TYPES, MergeError
 from xarray.core.options import OPTIONS, _get_keep_attrs
+from xarray.core.types import T_DataArrayOrSet
 from xarray.core.utils import (
     Default,
     HybridMappingProxy,
@@ -1844,7 +1845,7 @@ def _reindex_callback(
 
     def reindex_like(
         self: T_DataArray,
-        other: DataArray | Dataset,
+        other: T_Xarray,
         method: ReindexMethodOptions = None,
         tolerance: int | float | Iterable[int | float] | None = None,
         copy: bool = True,
@@ -2248,7 +2249,7 @@ def interp(
 
     def interp_like(
         self: T_DataArray,
-        other: DataArray | Dataset,
+        other: T_DataArrayOrSet,
         method: InterpOptions = "linear",
         assume_sorted: bool = False,
         kwargs: Mapping[str, Any] | None = None,
@@ -5375,7 +5376,7 @@ def map_blocks(
         func: Callable[..., T_Xarray],
         args: Sequence[Any] = (),
         kwargs: Mapping[str, Any] | None = None,
-        template: DataArray | Dataset | None = None,
+        template: T_Xarray | None = None,
     ) -> T_Xarray:
         """
         Apply a function to each block of this DataArray.

xarray/core/dataset.py

@@ -3398,7 +3398,7 @@ def _reindex_callback(
 
     def reindex_like(
         self: T_Dataset,
-        other: Dataset | DataArray,
+        other: T_Xarray,
         method: ReindexMethodOptions = None,
         tolerance: int | float | Iterable[int | float] | None = None,
         copy: bool = True,
@@ -8545,7 +8545,7 @@ def map_blocks(
         func: Callable[..., T_Xarray],
         args: Sequence[Any] = (),
         kwargs: Mapping[str, Any] | None = None,
-        template: DataArray | Dataset | None = None,
+        template: T_Xarray | None = None,
     ) -> T_Xarray:
         """
         Apply a function to each block of this Dataset.

xarray/core/parallel.py

@@ -30,9 +30,7 @@ def assert_chunks_compatible(a: Dataset, b: Dataset):
             raise ValueError(f"Chunk sizes along dimension {dim!r} are not equal.")
 
 
-def check_result_variables(
-    result: DataArray | Dataset, expected: Mapping[str, Any], kind: str
-):
+def check_result_variables(result: T_Xarray, expected: Mapping[str, Any], kind: str):
     if kind == "coords":
         nice_str = "coordinate"
     elif kind == "data_vars":
@@ -105,7 +103,7 @@ def make_meta(obj):
 
 
 def infer_template(
-    func: Callable[..., T_Xarray], obj: DataArray | Dataset, *args, **kwargs
+    func: Callable[..., T_Xarray], obj: T_Xarray, *args, **kwargs
 ) -> T_Xarray:
     """Infer return object by running the function on meta objects."""
     meta_args = [make_meta(arg) for arg in (obj,) + args]

xarray/plot/utils.py

@@ -14,6 +14,7 @@
 from xarray.core.indexes import PandasMultiIndex
 from xarray.core.options import OPTIONS
 from xarray.core.pycompat import DuckArrayModule
+from xarray.core.types import T_Xarray
 from xarray.core.utils import is_scalar, module_available
 
 nc_time_axis_available = module_available("nc_time_axis")
@@ -32,7 +33,6 @@
     from numpy.typing import ArrayLike
 
     from xarray.core.dataarray import DataArray
-    from xarray.core.dataset import Dataset
     from xarray.core.types import AspectOptions, ScaleOptions
 
     try:
@@ -312,7 +312,7 @@ def _determine_cmap_params(
 
 
 def _infer_xy_labels_3d(
-    darray: DataArray | Dataset,
+    darray: T_Xarray,
     x: Hashable | None,
     y: Hashable | None,
     rgb: Hashable | None,
@@ -374,7 +374,7 @@ def _infer_xy_labels_3d(
 
 
 def _infer_xy_labels(
-    darray: DataArray | Dataset,
+    darray: T_Xarray,
     x: Hashable | None,
     y: Hashable | None,
     imshow: bool = False,
@@ -413,9 +413,7 @@ def _infer_xy_labels(
 
 
 # TODO: Can by used to more than x or y, rename?
-def _assert_valid_xy(
-    darray: DataArray | Dataset, xy: Hashable | None, name: str
-) -> None:
+def _assert_valid_xy(darray: T_Xarray, xy: Hashable | None, name: str) -> None:
     """
     make sure x and y passed to plotting functions are valid
     """