diff --git a/xarray/core/computation.py b/xarray/core/computation.py index 6f1e08bf84f..823cbe02560 100644 --- a/xarray/core/computation.py +++ b/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 diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py index c59f1a6584f..737ed82bc05 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)