diff --git a/pandas/core/nanops.py b/pandas/core/nanops.py index d9e51810f1445..d38974839394d 100644 --- a/pandas/core/nanops.py +++ b/pandas/core/nanops.py @@ -367,14 +367,21 @@ def _wrap_results(result, dtype: np.dtype, fill_value=None): return result -def _datetimelike_compat(func): +def _datetimelike_compat(func: F) -> F: """ If we have datetime64 or timedelta64 values, ensure we have a correct mask before calling the wrapped function, then cast back afterwards. """ @functools.wraps(func) - def new_func(values, *, axis=None, skipna=True, mask=None, **kwargs): + def new_func( + values: np.ndarray, + *, + axis: Optional[int] = None, + skipna: bool = True, + mask: Optional[np.ndarray] = None, + **kwargs, + ): orig_values = values datetimelike = values.dtype.kind in ["m", "M"] @@ -390,7 +397,7 @@ def new_func(values, *, axis=None, skipna=True, mask=None, **kwargs): return result - return new_func + return cast(F, new_func) def _na_for_min_count(