From 5802801d4f53995a9ab21d94c6e5beedea94b2e8 Mon Sep 17 00:00:00 2001
From: Michael Niklas <mick.niklas@gmail.com>
Date: Wed, 11 May 2022 20:52:34 +0200
Subject: [PATCH 1/2] add polyval overloads

---
 xarray/core/computation.py       | 33 ++++++++++++++++++--------------
 xarray/tests/test_computation.py |  9 +++++++--
 2 files changed, 26 insertions(+), 16 deletions(-)

diff --git a/xarray/core/computation.py b/xarray/core/computation.py
index 6f1e08bf84f..faf313b2eee 100644
--- a/xarray/core/computation.py
+++ b/xarray/core/computation.py
@@ -16,7 +16,9 @@
     Hashable,
     Iterable,
     Mapping,
+    overload,
     Sequence,
+    TypeVar,
 )
 
 import numpy as np
@@ -1845,26 +1847,29 @@ 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 +1904,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
diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py
index c59f1a6584f..19efec65f48 100644
--- a/xarray/tests/test_computation.py
+++ b/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)
 
 

From a67523c21a077171759e6dc9808e50b2376c1cc3 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 11 May 2022 18:56:31 +0000
Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 xarray/core/computation.py       | 9 +++++----
 xarray/tests/test_computation.py | 2 +-
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/xarray/core/computation.py b/xarray/core/computation.py
index faf313b2eee..823cbe02560 100644
--- a/xarray/core/computation.py
+++ b/xarray/core/computation.py
@@ -16,9 +16,8 @@
     Hashable,
     Iterable,
     Mapping,
-    overload,
     Sequence,
-    TypeVar,
+    overload,
 )
 
 import numpy as np
@@ -1861,14 +1860,16 @@ def polyval(coord: DataArray, coeffs: Dataset, degree_dim: Hashable) -> Dataset:
 def polyval(coord: Dataset, coeffs: DataArray, degree_dim: Hashable) -> Dataset:
     ...
 
-    
+
 @overload
 def polyval(coord: Dataset, coeffs: Dataset, degree_dim: Hashable) -> Dataset:
     ...
 
 
 def polyval(
-    coord: Dataset | DataArray, coeffs: Dataset | DataArray, degree_dim: Hashable = "degree"
+    coord: Dataset | DataArray,
+    coeffs: Dataset | DataArray,
+    degree_dim: Hashable = "degree",
 ) -> Dataset | DataArray:
     """Evaluate a polynomial at specific values
 
diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py
index 19efec65f48..737ed82bc05 100644
--- a/xarray/tests/test_computation.py
+++ b/xarray/tests/test_computation.py
@@ -2016,7 +2016,7 @@ def test_polyval(
     use_dask: bool,
     x: xr.DataArray | xr.Dataset,
     coeffs: xr.DataArray | xr.Dataset,
-    expected: xr.DataArray | xr.Dataset
+    expected: xr.DataArray | xr.Dataset,
 ) -> None:
     if use_dask:
         if not has_dask: