diff --git a/pandas/core/nanops.py b/pandas/core/nanops.py index e775156a6ae2f..1c35eacf55c03 100644 --- a/pandas/core/nanops.py +++ b/pandas/core/nanops.py @@ -726,7 +726,7 @@ def nanmean( @bottleneck_switch() -def nanmedian(values, *, axis: AxisInt | None = None, skipna: bool = True, mask=None): +def nanmedian(values: np.ndarray, *, axis: AxisInt | None = None, skipna: bool = True, mask=None) -> float | np.ndarray: """ Parameters ---------- @@ -738,7 +738,7 @@ def nanmedian(values, *, axis: AxisInt | None = None, skipna: bool = True, mask= Returns ------- - result : float + result : float | ndarray Unless input is a float array, in which case use the same precision as the input array. @@ -758,7 +758,7 @@ def nanmedian(values, *, axis: AxisInt | None = None, skipna: bool = True, mask= # cases we never need to set NaN to the masked values using_nan_sentinel = values.dtype.kind == "f" and mask is None - def get_median(x, _mask=None): + def get_median(x: np.ndarray, _mask=None): if _mask is None: _mask = notna(x) else: @@ -794,6 +794,8 @@ def get_median(x, _mask=None): notempty = values.size + res: float | np.ndarray + # an array from a frame if values.ndim > 1 and axis is not None: # there's a non-empty array to apply over otherwise numpy raises