Improve alignment typehints (#4522)

* Improve alignment typehints

* Fix typing issues

* Add a note to what's new

Author: m1so

doc/whats-new.rst

@@ -36,6 +36,8 @@ New Features
   limited regions of existing Zarr stores (:pull:`4035`).
   See :ref:`io.zarr.appending` for full details.
   By `Stephan Hoyer <https://github.com/shoyer>`_.
+- Added typehints in :py:func:`align` to reflect that the same type received in ``objects`` arg will be returned (:pull:`4522`).
+  By `Michal Baumgartner <https://github.com/m1so>`_.
 
 Bug fixes
 ~~~~~~~~~

xarray/core/alignment.py

@@ -2,7 +2,17 @@
 import operator
 from collections import defaultdict
 from contextlib import suppress
-from typing import TYPE_CHECKING, Any, Dict, Hashable, Mapping, Optional, Tuple, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Hashable,
+    Mapping,
+    Optional,
+    Tuple,
+    TypeVar,
+    Union,
+)
 
 import numpy as np
 import pandas as pd
@@ -13,9 +23,12 @@
 from .variable import IndexVariable, Variable
 
 if TYPE_CHECKING:
+    from .common import DataWithCoords
     from .dataarray import DataArray
     from .dataset import Dataset
 
+    DataAlignable = TypeVar("DataAlignable", bound=DataWithCoords)
+
 
 def _get_joiner(join):
     if join == "outer":
@@ -59,13 +72,13 @@ def _override_indexes(objects, all_indexes, exclude):
 
 
 def align(
-    *objects,
+    *objects: "DataAlignable",
     join="inner",
     copy=True,
     indexes=None,
     exclude=frozenset(),
     fill_value=dtypes.NA,
-):
+) -> Tuple["DataAlignable", ...]:
     """
     Given any number of Dataset and/or DataArray objects, returns new
     objects with aligned indexes and dimension sizes.
@@ -337,7 +350,9 @@ def align(
             # fast path for no reindexing necessary
             new_obj = obj.copy(deep=copy)
         else:
-            new_obj = obj.reindex(copy=copy, fill_value=fill_value, **valid_indexers)
+            new_obj = obj.reindex(
+                copy=copy, fill_value=fill_value, indexers=valid_indexers
+            )
         new_obj.encoding = obj.encoding
         result.append(new_obj)
 

xarray/core/concat.py

@@ -380,8 +380,8 @@ def _dataset_concat(
     dim, coord = _calc_concat_dim_coord(dim)
     # Make sure we're working on a copy (we'll be loading variables)
     datasets = [ds.copy() for ds in datasets]
-    datasets = align(
-        *datasets, join=join, copy=False, exclude=[dim], fill_value=fill_value
+    datasets = list(
+        align(*datasets, join=join, copy=False, exclude=[dim], fill_value=fill_value)
     )
 
     dim_coords, dims_sizes, coord_names, data_names = _parse_datasets(datasets)