Fix polyval overloads (#6593)

* add polyval overloads

* [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>

Author: headtr1ck

xarray/core/computation.py

@@ -17,6 +17,7 @@
     Iterable,
     Mapping,
     Sequence,
+    overload,
 )
 
 import numpy as np
@@ -1845,26 +1846,31 @@ def where(cond, x, y, keep_attrs=None):
     )
 
 
-# These overloads seem not to work — mypy says it can't find a matching overload for
-# `DataArray` & `DataArray`, despite that being in the first overload. Would be nice to
-# have overloaded functions rather than just `T_Xarray` for everything.
+@overload
+def polyval(coord: DataArray, coeffs: DataArray, degree_dim: Hashable) -> DataArray:
+    ...
 
-# @overload
-# def polyval(coord: DataArray, coeffs: DataArray, degree_dim: Hashable) -> DataArray:
-#     ...
 
+@overload
+def polyval(coord: DataArray, coeffs: Dataset, degree_dim: Hashable) -> Dataset:
+    ...
 
-# @overload
-# def polyval(coord: T_Xarray, coeffs: Dataset, degree_dim: Hashable) -> Dataset:
-#     ...
 
+@overload
+def polyval(coord: Dataset, coeffs: DataArray, degree_dim: Hashable) -> Dataset:
+    ...
 
-# @overload
-# def polyval(coord: Dataset, coeffs: T_Xarray, degree_dim: Hashable) -> Dataset:
-#     ...
+
+@overload
+def polyval(coord: Dataset, coeffs: Dataset, degree_dim: Hashable) -> Dataset:
+    ...
 
 
-def polyval(coord: T_Xarray, coeffs: T_Xarray, degree_dim="degree") -> T_Xarray:
+def polyval(
+    coord: Dataset | DataArray,
+    coeffs: Dataset | DataArray,
+    degree_dim: Hashable = "degree",
+) -> Dataset | DataArray:
     """Evaluate a polynomial at specific values
 
     Parameters
@@ -1899,7 +1905,7 @@ def polyval(coord: T_Xarray, coeffs: T_Xarray, degree_dim="degree") -> T_Xarray:
     coeffs = coeffs.reindex(
         {degree_dim: np.arange(max_deg + 1)}, fill_value=0, copy=False
     )
-    coord = _ensure_numeric(coord)
+    coord = _ensure_numeric(coord)  # type: ignore # https://github.com/python/mypy/issues/1533 ?
 
     # using Horner's method
     # https://en.wikipedia.org/wiki/Horner%27s_method

xarray/tests/test_computation.py

@@ -2012,14 +2012,19 @@ def test_where_attrs() -> None:
         ),
     ],
 )
-def test_polyval(use_dask, x: T_Xarray, coeffs: T_Xarray, expected) -> None:
+def test_polyval(
+    use_dask: bool,
+    x: xr.DataArray | xr.Dataset,
+    coeffs: xr.DataArray | xr.Dataset,
+    expected: xr.DataArray | xr.Dataset,
+) -> None:
     if use_dask:
         if not has_dask:
             pytest.skip("requires dask")
         coeffs = coeffs.chunk({"degree": 2})
         x = x.chunk({"x": 2})
     with raise_if_dask_computes():
-        actual = xr.polyval(coord=x, coeffs=coeffs)
+        actual = xr.polyval(coord=x, coeffs=coeffs)  # type: ignore
     xr.testing.assert_allclose(actual, expected)