Improved typing of align & broadcast (#8234)

* add overloads to align

* add overloads to broadcast as well

* add some more typing

* remove unused ignore

Author: headtr1ck

xarray/core/alignment.py

@@ -5,7 +5,7 @@
 from collections import defaultdict
 from collections.abc import Hashable, Iterable, Mapping
 from contextlib import suppress
-from typing import TYPE_CHECKING, Any, Callable, Generic, cast
+from typing import TYPE_CHECKING, Any, Callable, Final, Generic, TypeVar, cast, overload
 
 import numpy as np
 import pandas as pd
@@ -26,7 +26,13 @@
 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
+    from xarray.core.types import (
+        Alignable,
+        JoinOptions,
+        T_DataArray,
+        T_Dataset,
+        T_DuckArray,
+    )
 
 
 def reindex_variables(
@@ -128,7 +134,7 @@ def __init__(
         objects: Iterable[T_Alignable],
         join: str = "inner",
         indexes: Mapping[Any, Any] | None = None,
-        exclude_dims: Iterable = frozenset(),
+        exclude_dims: str | Iterable[Hashable] = frozenset(),
         exclude_vars: Iterable[Hashable] = frozenset(),
         method: str | None = None,
         tolerance: int | float | Iterable[int | float] | None = None,
@@ -576,12 +582,111 @@ def align(self) -> None:
             self.reindex_all()
 
 
+T_Obj1 = TypeVar("T_Obj1", bound="Alignable")
+T_Obj2 = TypeVar("T_Obj2", bound="Alignable")
+T_Obj3 = TypeVar("T_Obj3", bound="Alignable")
+T_Obj4 = TypeVar("T_Obj4", bound="Alignable")
+T_Obj5 = TypeVar("T_Obj5", bound="Alignable")
+
+
+@overload
+def align(
+    obj1: T_Obj1,
+    /,
+    *,
+    join: JoinOptions = "inner",
+    copy: bool = True,
+    indexes=None,
+    exclude: str | Iterable[Hashable] = frozenset(),
+    fill_value=dtypes.NA,
+) -> tuple[T_Obj1]:
+    ...
+
+
+@overload
+def align(  # type: ignore[misc]
+    obj1: T_Obj1,
+    obj2: T_Obj2,
+    /,
+    *,
+    join: JoinOptions = "inner",
+    copy: bool = True,
+    indexes=None,
+    exclude: str | Iterable[Hashable] = frozenset(),
+    fill_value=dtypes.NA,
+) -> tuple[T_Obj1, T_Obj2]:
+    ...
+
+
+@overload
+def align(  # type: ignore[misc]
+    obj1: T_Obj1,
+    obj2: T_Obj2,
+    obj3: T_Obj3,
+    /,
+    *,
+    join: JoinOptions = "inner",
+    copy: bool = True,
+    indexes=None,
+    exclude: str | Iterable[Hashable] = frozenset(),
+    fill_value=dtypes.NA,
+) -> tuple[T_Obj1, T_Obj2, T_Obj3]:
+    ...
+
+
+@overload
+def align(  # type: ignore[misc]
+    obj1: T_Obj1,
+    obj2: T_Obj2,
+    obj3: T_Obj3,
+    obj4: T_Obj4,
+    /,
+    *,
+    join: JoinOptions = "inner",
+    copy: bool = True,
+    indexes=None,
+    exclude: str | Iterable[Hashable] = frozenset(),
+    fill_value=dtypes.NA,
+) -> tuple[T_Obj1, T_Obj2, T_Obj3, T_Obj4]:
+    ...
+
+
+@overload
+def align(  # type: ignore[misc]
+    obj1: T_Obj1,
+    obj2: T_Obj2,
+    obj3: T_Obj3,
+    obj4: T_Obj4,
+    obj5: T_Obj5,
+    /,
+    *,
+    join: JoinOptions = "inner",
+    copy: bool = True,
+    indexes=None,
+    exclude: str | Iterable[Hashable] = frozenset(),
+    fill_value=dtypes.NA,
+) -> tuple[T_Obj1, T_Obj2, T_Obj3, T_Obj4, T_Obj5]:
+    ...
+
+
+@overload
 def align(
     *objects: T_Alignable,
     join: JoinOptions = "inner",
     copy: bool = True,
     indexes=None,
-    exclude=frozenset(),
+    exclude: str | Iterable[Hashable] = frozenset(),
+    fill_value=dtypes.NA,
+) -> tuple[T_Alignable, ...]:
+    ...
+
+
+def align(  # type: ignore[misc]
+    *objects: T_Alignable,
+    join: JoinOptions = "inner",
+    copy: bool = True,
+    indexes=None,
+    exclude: str | Iterable[Hashable] = frozenset(),
     fill_value=dtypes.NA,
 ) -> tuple[T_Alignable, ...]:
     """
@@ -620,7 +725,7 @@ def align(
     indexes : dict-like, optional
         Any indexes explicitly provided with the `indexes` argument should be
         used in preference to the aligned indexes.
-    exclude : sequence of str, optional
+    exclude : str, iterable of hashable or None, optional
         Dimensions that must be excluded from alignment
     fill_value : scalar or dict-like, optional
         Value to use for newly missing values. If a dict-like, maps
@@ -787,12 +892,12 @@ def align(
 def deep_align(
     objects: Iterable[Any],
     join: JoinOptions = "inner",
-    copy=True,
+    copy: bool = True,
     indexes=None,
-    exclude=frozenset(),
-    raise_on_invalid=True,
+    exclude: str | Iterable[Hashable] = frozenset(),
+    raise_on_invalid: bool = True,
     fill_value=dtypes.NA,
-):
+) -> list[Any]:
     """Align objects for merging, recursing into dictionary values.
 
     This function is not public API.
@@ -807,12 +912,12 @@ def deep_align(
     def is_alignable(obj):
         return isinstance(obj, (Coordinates, DataArray, Dataset))
 
-    positions = []
-    keys = []
-    out = []
-    targets = []
-    no_key = object()
-    not_replaced = object()
+    positions: list[int] = []
+    keys: list[type[object] | Hashable] = []
+    out: list[Any] = []
+    targets: list[Alignable] = []
+    no_key: Final = object()
+    not_replaced: Final = object()
     for position, variables in enumerate(objects):
         if is_alignable(variables):
             positions.append(position)
@@ -857,7 +962,7 @@ def is_alignable(obj):
         if key is no_key:
             out[position] = aligned_obj
         else:
-            out[position][key] = aligned_obj  # type: ignore[index]  # maybe someone can fix this?
+            out[position][key] = aligned_obj
 
     return out
 
@@ -988,9 +1093,69 @@ def _broadcast_dataset(ds: T_Dataset) -> T_Dataset:
         raise ValueError("all input must be Dataset or DataArray objects")
 
 
-# TODO: this typing is too restrictive since it cannot deal with mixed
-# DataArray and Dataset types...? Is this a problem?
-def broadcast(*args: T_Alignable, exclude=None) -> tuple[T_Alignable, ...]:
+@overload
+def broadcast(
+    obj1: T_Obj1, /, *, exclude: str | Iterable[Hashable] | None = None
+) -> tuple[T_Obj1]:
+    ...
+
+
+@overload
+def broadcast(  # type: ignore[misc]
+    obj1: T_Obj1, obj2: T_Obj2, /, *, exclude: str | Iterable[Hashable] | None = None
+) -> tuple[T_Obj1, T_Obj2]:
+    ...
+
+
+@overload
+def broadcast(  # type: ignore[misc]
+    obj1: T_Obj1,
+    obj2: T_Obj2,
+    obj3: T_Obj3,
+    /,
+    *,
+    exclude: str | Iterable[Hashable] | None = None,
+) -> tuple[T_Obj1, T_Obj2, T_Obj3]:
+    ...
+
+
+@overload
+def broadcast(  # type: ignore[misc]
+    obj1: T_Obj1,
+    obj2: T_Obj2,
+    obj3: T_Obj3,
+    obj4: T_Obj4,
+    /,
+    *,
+    exclude: str | Iterable[Hashable] | None = None,
+) -> tuple[T_Obj1, T_Obj2, T_Obj3, T_Obj4]:
+    ...
+
+
+@overload
+def broadcast(  # type: ignore[misc]
+    obj1: T_Obj1,
+    obj2: T_Obj2,
+    obj3: T_Obj3,
+    obj4: T_Obj4,
+    obj5: T_Obj5,
+    /,
+    *,
+    exclude: str | Iterable[Hashable] | None = None,
+) -> tuple[T_Obj1, T_Obj2, T_Obj3, T_Obj4, T_Obj5]:
+    ...
+
+
+@overload
+def broadcast(
+    *args: T_Alignable, exclude: str | Iterable[Hashable] | None = None
+) -> tuple[T_Alignable, ...]:
+    ...
+
+
+def broadcast(  # type: ignore[misc]
+    *args: T_Alignable, exclude: str | Iterable[Hashable] | None = None
+) -> tuple[T_Alignable, ...]:
     """Explicitly broadcast any number of DataArray or Dataset objects against
     one another.
 
@@ -1004,7 +1169,7 @@ def broadcast(*args: T_Alignable, exclude=None) -> tuple[T_Alignable, ...]:
     ----------
     *args : DataArray or Dataset
         Arrays to broadcast against each other.
-    exclude : sequence of str, optional
+    exclude : str, iterable of hashable or None, optional
         Dimensions that must not be broadcasted
 
     Returns

xarray/core/common.py

@@ -1163,7 +1163,7 @@ def where(self, cond: Any, other: Any = dtypes.NA, drop: bool = False) -> Self:
                     f"cond argument is {cond!r} but must be a {Dataset!r} or {DataArray!r} (or a callable than returns one)."
                 )
 
-            self, cond = align(self, cond)  # type: ignore[assignment]
+            self, cond = align(self, cond)
 
             def _dataarray_indexer(dim: Hashable) -> DataArray:
                 return cond.any(dim=(d for d in cond.dims if d != dim))

xarray/core/computation.py

@@ -289,8 +289,14 @@ def apply_dataarray_vfunc(
     from xarray.core.dataarray import DataArray
 
     if len(args) > 1:
-        args = deep_align(
-            args, join=join, copy=False, exclude=exclude_dims, raise_on_invalid=False
+        args = tuple(
+            deep_align(
+                args,
+                join=join,
+                copy=False,
+                exclude=exclude_dims,
+                raise_on_invalid=False,
+            )
         )
 
     objs = _all_of_type(args, DataArray)
@@ -506,8 +512,14 @@ def apply_dataset_vfunc(
     objs = _all_of_type(args, Dataset)
 
     if len(args) > 1:
-        args = deep_align(
-            args, join=join, copy=False, exclude=exclude_dims, raise_on_invalid=False
+        args = tuple(
+            deep_align(
+                args,
+                join=join,
+                copy=False,
+                exclude=exclude_dims,
+                raise_on_invalid=False,
+            )
         )
 
     list_of_coords, list_of_indexes = build_output_coords_and_indexes(

xarray/core/dataarray.py

@@ -4639,7 +4639,7 @@ def _binary_op(
             return NotImplemented
         if isinstance(other, DataArray):
             align_type = OPTIONS["arithmetic_join"]
-            self, other = align(self, other, join=align_type, copy=False)  # type: ignore[type-var,assignment]
+            self, other = align(self, other, join=align_type, copy=False)
         other_variable_or_arraylike: DaCompatible = getattr(other, "variable", other)
         other_coords = getattr(other, "coords", None)
 

xarray/core/dataset.py

@@ -7508,7 +7508,7 @@ def _binary_op(self, other, f, reflexive=False, join=None) -> Dataset:
             return NotImplemented
         align_type = OPTIONS["arithmetic_join"] if join is None else join
         if isinstance(other, (DataArray, Dataset)):
-            self, other = align(self, other, join=align_type, copy=False)  # type: ignore[assignment]
+            self, other = align(self, other, join=align_type, copy=False)
         g = f if not reflexive else lambda x, y: f(y, x)
         ds = self._calculate_binary_op(g, other, join=align_type)
         keep_attrs = _get_keep_attrs(default=False)
@@ -7920,9 +7920,9 @@ def sortby(
         else:
             variables = variables
         arrays = [v if isinstance(v, DataArray) else self[v] for v in variables]
-        aligned_vars = align(self, *arrays, join="left")  # type: ignore[type-var]
-        aligned_self = cast(Self, aligned_vars[0])
-        aligned_other_vars: tuple[DataArray, ...] = aligned_vars[1:]  # type: ignore[assignment]
+        aligned_vars = align(self, *arrays, join="left")
+        aligned_self = aligned_vars[0]
+        aligned_other_vars: tuple[DataArray, ...] = aligned_vars[1:]
         vars_by_dim = defaultdict(list)
         for data_array in aligned_other_vars:
             if data_array.ndim != 1:

xarray/core/merge.py

@@ -474,10 +474,11 @@ def coerce_pandas_values(objects: Iterable[CoercibleMapping]) -> list[DatasetLik
     from xarray.core.dataarray import DataArray
     from xarray.core.dataset import Dataset
 
-    out = []
+    out: list[DatasetLike] = []
     for obj in objects:
+        variables: DatasetLike
         if isinstance(obj, (Dataset, Coordinates)):
-            variables: DatasetLike = obj
+            variables = obj
         else:
             variables = {}
             if isinstance(obj, PANDAS_TYPES):
@@ -491,7 +492,7 @@ def coerce_pandas_values(objects: Iterable[CoercibleMapping]) -> list[DatasetLik
 
 
 def _get_priority_vars_and_indexes(
-    objects: list[DatasetLike],
+    objects: Sequence[DatasetLike],
     priority_arg: int | None,
     compat: CompatOptions = "equals",
 ) -> dict[Hashable, MergeElement]:
@@ -503,7 +504,7 @@ def _get_priority_vars_and_indexes(
 
     Parameters
     ----------
-    objects : list of dict-like of Variable
+    objects : sequence of dict-like of Variable
         Dictionaries in which to find the priority variables.
     priority_arg : int or None
         Integer object whose variable should take priority.