diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index b5bcc255b70..0341d84b661 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -87,7 +87,6 @@ DatetimeUnitOptions, ErrorOptions, ErrorOptionsWithWarn, - InterpAllOptions, InterpOptions, PadModeOptions, PadReflectOptions, @@ -2626,7 +2625,7 @@ def fillna(self: T_DataArray, value: Any) -> T_DataArray: def interpolate_na( self: T_DataArray, dim: Hashable | None = None, - method: InterpAllOptions = "linear", + method: InterpOptions = "linear", limit: int | None = None, use_coordinate: bool | str = True, max_gap: ( diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index c5c727f4bed..084437e102c 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -109,6 +109,7 @@ CompatOptions, ErrorOptions, ErrorOptionsWithWarn, + InterpOptions, JoinOptions, PadModeOptions, PadReflectOptions, @@ -3037,7 +3038,7 @@ def _reindex( def interp( self, coords: Mapping[Any, Any] = None, - method: str = "linear", + method: InterpOptions = "linear", assume_sorted: bool = False, kwargs: Mapping[str, Any] = None, method_non_numeric: str = "nearest", @@ -3298,7 +3299,7 @@ def _validate_interp_indexer(x, new_x): def interp_like( self, other: Dataset | DataArray, - method: str = "linear", + method: InterpOptions = "linear", assume_sorted: bool = False, kwargs: Mapping[str, Any] = None, method_non_numeric: str = "nearest", diff --git a/xarray/core/missing.py b/xarray/core/missing.py index 2e869dbe675..5e954c8ce27 100644 --- a/xarray/core/missing.py +++ b/xarray/core/missing.py @@ -4,7 +4,7 @@ import warnings from functools import partial from numbers import Number -from typing import TYPE_CHECKING, Any, Callable, Hashable, Sequence +from typing import TYPE_CHECKING, Any, Callable, Hashable, Sequence, get_args import numpy as np import pandas as pd @@ -16,6 +16,7 @@ from .duck_array_ops import datetime_to_numeric, push, timedelta_to_numeric from .options import OPTIONS, _get_keep_attrs from .pycompat import dask_version, is_duck_dask_array +from .types import Interp1dOptions, InterpOptions from .utils import OrderedSet, is_scalar from .variable import Variable, broadcast_variables @@ -309,7 +310,7 @@ def interp_na( self, dim: Hashable = None, use_coordinate: bool | str = True, - method: str = "linear", + method: InterpOptions = "linear", limit: int = None, max_gap: int | float | str | pd.Timedelta | np.timedelta64 | dt.timedelta = None, keep_attrs: bool = None, @@ -469,28 +470,20 @@ def _import_interpolant(interpolant, method): raise ImportError(f"Interpolation with method {method} requires scipy.") from e -def _get_interpolator(method, vectorizeable_only=False, **kwargs): +def _get_interpolator( + method: InterpOptions, vectorizeable_only: bool = False, **kwargs +): """helper function to select the appropriate interpolator class returns interpolator class and keyword arguments for the class """ - interp1d_methods = [ - "linear", - "nearest", - "zero", - "slinear", - "quadratic", - "cubic", - "polynomial", - ] - valid_methods = interp1d_methods + [ - "barycentric", - "krog", - "pchip", - "spline", - "akima", + interp_class: type[NumpyInterpolator] | type[ScipyInterpolator] | type[ + SplineInterpolator ] + interp1d_methods = get_args(Interp1dOptions) + valid_methods = tuple(vv for v in get_args(InterpOptions) for vv in get_args(v)) + # prioritize scipy.interpolate if ( method == "linear" @@ -597,7 +590,7 @@ def _floatize_x(x, new_x): return x, new_x -def interp(var, indexes_coords, method, **kwargs): +def interp(var, indexes_coords, method: InterpOptions, **kwargs): """Make an interpolation of Variable Parameters @@ -650,7 +643,7 @@ def interp(var, indexes_coords, method, **kwargs): result = Variable(new_dims, interped, attrs=var.attrs) # dimension of the output array - out_dims = OrderedSet() + out_dims: OrderedSet = OrderedSet() for d in var.dims: if d in dims: out_dims.update(indexes_coords[d][1].dims) @@ -660,7 +653,7 @@ def interp(var, indexes_coords, method, **kwargs): return result -def interp_func(var, x, new_x, method, kwargs): +def interp_func(var, x, new_x, method: InterpOptions, kwargs): """ multi-dimensional interpolation for array-like. Interpolated axes should be located in the last position. diff --git a/xarray/core/types.py b/xarray/core/types.py index 38d13c07c3c..f4f86bafc93 100644 --- a/xarray/core/types.py +++ b/xarray/core/types.py @@ -49,11 +49,11 @@ ] JoinOptions = Literal["outer", "inner", "left", "right", "exact", "override"] -InterpOptions = Literal["linear", "nearest", "zero", "slinear", "quadratic", "cubic"] -Interp1dOptions = Union[InterpOptions, Literal["polynomial"]] -InterpAllOptions = Union[ - Interp1dOptions, Literal["barycentric", "krog", "pchip", "spline", "akima"] +Interp1dOptions = Literal[ + "linear", "nearest", "zero", "slinear", "quadratic", "cubic", "polynomial" ] +InterpolantOptions = Literal["barycentric", "krog", "pchip", "spline", "akima"] +InterpOptions = Union[Interp1dOptions, InterpolantOptions] DatetimeUnitOptions = Literal[ "Y", "M", "W", "D", "h", "m", "s", "ms", "us", "ns", "ps", "fs", "as"