Add typing to functions related to data_vars (#8226)

* Update dataset.py * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * more typing * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update xarray/core/dataset.py Co-authored-by: Michael Niklas <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Michael Niklas <[email protected]>
pydata · Sep 24, 2023 · 77eaa8b · 77eaa8b
1 parent b14fbd9
commit 77eaa8b
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 9 deletions.
diff --git a/xarray/core/coordinates.py b/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).
 

diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py
@@ -124,7 +124,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 +133,7 @@
         CoarsenBoundaryOptions,
         CombineAttrsOptions,
         CompatOptions,
+        DataVars,
         DatetimeLike,
         DatetimeUnitOptions,
         Dims,
@@ -404,7 +405,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 +667,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 +1221,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 +1323,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:

diff --git a/xarray/core/types.py b/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"]