Add typing to functions related to data_vars

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, T_Xarray
+from xarray.core.types import DataVars, Self, T_DataArray, T_Xarray
 from xarray.core.utils import (
     Frozen,
     ReprObject,
@@ -937,7 +937,7 @@ def assert_coordinate_consistent(obj: T_Xarray, coords: Mapping[Any, Variable])
 
 
 def create_coords_with_default_indexes(
-    coords: Mapping[Any, Any], data_vars: Mapping[Any, Any] | None = None
+    coords: Mapping[Any, Any], data_vars: DataVars | None = None
 ) -> Coordinates:
     """Returns a Coordinates object from a mapping of coordinates (arbitrary objects).
 

xarray/core/dataset.py

@@ -21,7 +21,16 @@
 from numbers import Number
 from operator import methodcaller
 from os import PathLike
-from typing import IO, TYPE_CHECKING, Any, Callable, Generic, Literal, cast, overload
+from typing import (
+    IO,
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Generic,
+    Literal,
+    cast,
+    overload,
+)
 
 import numpy as np
 
@@ -124,7 +133,7 @@
     from xarray.backends.api import T_NetcdfEngine, T_NetcdfTypes
     from xarray.core.dataarray import DataArray
     from xarray.core.groupby import DatasetGroupBy
-    from xarray.core.merge import CoercibleMapping, CoercibleValue
+    from xarray.core.merge import CoercibleMapping, CoercibleValue, _MergeResult
     from xarray.core.parallelcompat import ChunkManagerEntrypoint
     from xarray.core.resample import DatasetResample
     from xarray.core.rolling import DatasetCoarsen, DatasetRolling
@@ -133,6 +142,7 @@
         CoarsenBoundaryOptions,
         CombineAttrsOptions,
         CompatOptions,
+        DataVars,
         DatetimeLike,
         DatetimeUnitOptions,
         Dims,
@@ -404,7 +414,7 @@ def _initialize_feasible(lb, ub):
     return param_defaults, bounds_defaults
 
 
-def merge_data_and_coords(data_vars, coords):
+def merge_data_and_coords(data_vars: DataVars, coords) -> _MergeResult:
     """Used in Dataset.__init__."""
     if isinstance(coords, Coordinates):
         coords = coords.copy()
@@ -666,7 +676,7 @@ def __init__(
         self,
         # could make a VariableArgs to use more generally, and refine these
         # categories
-        data_vars: Mapping[Any, Any] | None = None,
+        data_vars: DataVars | None = None,
         coords: Mapping[Any, Any] | None = None,
         attrs: Mapping[Any, Any] | None = None,
     ) -> None:
@@ -1220,9 +1230,7 @@ def _overwrite_indexes(
         else:
             return replaced
 
-    def copy(
-        self, deep: bool = False, data: Mapping[Any, ArrayLike] | None = None
-    ) -> Self:
+    def copy(self, deep: bool = False, data: DataVars | None = None) -> Self:
         """Returns a copy of this dataset.
 
         If `deep=True`, a deep copy is made of each of the component variables.
@@ -1324,7 +1332,7 @@ def copy(
     def _copy(
         self,
         deep: bool = False,
-        data: Mapping[Any, ArrayLike] | None = None,
+        data: DataVars | None = None,
         memo: dict[int, Any] | None = None,
     ) -> Self:
         if data is None:

xarray/core/types.py

@@ -187,6 +187,9 @@ def copy(
 T_Chunks = Union[int, dict[Any, Any], Literal["auto"], None]
 T_NormalizedChunks = tuple[tuple[int, ...], ...]
 
+DataVars = Mapping[Any, Any]
+
+
 ErrorOptions = Literal["raise", "ignore"]
 ErrorOptionsWithWarn = Literal["raise", "warn", "ignore"]