TYP: Annotations in pandas/core/nanops.py

Author: MomIsBestFriend

pandas/core/nanops.py

@@ -1,13 +1,14 @@
 import functools
 import itertools
 import operator
-from typing import Any, Optional, Tuple, Union
+from typing import Any, Callable, Optional, Tuple, Union
 
 import numpy as np
 
 from pandas._config import get_option
 
 from pandas._libs import NaT, Timedelta, Timestamp, iNaT, lib
+from pandas._typing import Dtype, Scalar
 from pandas.compat._optional import import_optional_dependency
 
 from pandas.core.dtypes.cast import _int64_max, maybe_upcast_putmask
@@ -37,7 +38,7 @@
 _USE_BOTTLENECK = False
 
 
-def set_use_bottleneck(v=True):
+def set_use_bottleneck(v: bool = True) -> None:
     # set/unset to use bottleneck
     global _USE_BOTTLENECK
     if _BOTTLENECK_INSTALLED:
@@ -55,7 +56,7 @@ def __init__(self, *dtypes):
     def check(self, obj) -> bool:
         return hasattr(obj, "dtype") and issubclass(obj.dtype.type, self.dtypes)
 
-    def __call__(self, f):
+    def __call__(self, f) -> Callable:
         @functools.wraps(f)
         def _f(*args, **kwargs):
             obj_iter = itertools.chain(args, kwargs.values())
@@ -80,11 +81,11 @@ def _f(*args, **kwargs):
 
 
 class bottleneck_switch:
-    def __init__(self, name=None, **kwargs):
+    def __init__(self, name: Optional[str] = None, **kwargs):
         self.name = name
         self.kwargs = kwargs
 
-    def __call__(self, alt):
+    def __call__(self, alt: Callable) -> Callable:
         bn_name = self.name or alt.__name__
 
         try:
@@ -93,7 +94,9 @@ def __call__(self, alt):
             bn_func = None
 
         @functools.wraps(alt)
-        def f(values, axis=None, skipna=True, **kwds):
+        def f(
+            values: np.ndarray, axis: Optional[int] = None, skipna: bool = True, **kwds
+        ):
             if len(self.kwargs) > 0:
                 for k, v in self.kwargs.items():
                     if k not in kwds:
@@ -129,7 +132,7 @@ def f(values, axis=None, skipna=True, **kwds):
         return f
 
 
-def _bn_ok_dtype(dt, name: str) -> bool:
+def _bn_ok_dtype(dt: Dtype, name: str) -> bool:
     # Bottleneck chokes on datetime64
     if not is_object_dtype(dt) and not (
         is_datetime_or_timedelta_dtype(dt) or is_datetime64tz_dtype(dt)
@@ -163,7 +166,9 @@ def _has_infs(result) -> bool:
         return False
 
 
-def _get_fill_value(dtype, fill_value=None, fill_value_typ=None):
+def _get_fill_value(
+    dtype: Dtype, fill_value: Any = None, fill_value_typ: Optional[str] = None
+):
     """ return the correct fill value for the dtype of the values """
     if fill_value is not None:
         return fill_value
@@ -326,12 +331,12 @@ def _get_values(
     return values, mask, dtype, dtype_max, fill_value
 
 
-def _na_ok_dtype(dtype):
+def _na_ok_dtype(dtype) -> bool:
     # TODO: what about datetime64tz?  PeriodDtype?
     return not issubclass(dtype.type, (np.integer, np.timedelta64, np.datetime64))
 
 
-def _wrap_results(result, dtype, fill_value=None):
+def _wrap_results(result, dtype: Dtype, fill_value=None):
     """ wrap our results if needed """
 
     if is_datetime64_dtype(dtype) or is_datetime64tz_dtype(dtype):
@@ -362,7 +367,9 @@ def _wrap_results(result, dtype, fill_value=None):
     return result
 
 
-def _na_for_min_count(values, axis: Optional[int]):
+def _na_for_min_count(
+    values: np.ndarray, axis: Optional[int]
+) -> Union[Scalar, np.ndarray]:
     """
     Return the missing value for `values`.
 
@@ -393,7 +400,12 @@ def _na_for_min_count(values, axis: Optional[int]):
         return result
 
 
-def nanany(values, axis=None, skipna: bool = True, mask=None):
+def nanany(
+    values: np.ndarray,
+    axis: Optional[int] = None,
+    skipna: bool = True,
+    mask: Optional[np.ndarray] = None,
+) -> bool:
     """
     Check if any elements along an axis evaluate to True.
 
@@ -425,7 +437,12 @@ def nanany(values, axis=None, skipna: bool = True, mask=None):
     return values.any(axis)
 
 
-def nanall(values, axis=None, skipna: bool = True, mask=None):
+def nanall(
+    values: np.ndarray,
+    axis: Optional[int] = None,
+    skipna: bool = True,
+    mask: Optional[np.ndarray] = None,
+) -> bool:
     """
     Check if all elements along an axis evaluate to True.
 
@@ -458,7 +475,13 @@ def nanall(values, axis=None, skipna: bool = True, mask=None):
 
 
 @disallow("M8")
-def nansum(values, axis=None, skipna=True, min_count=0, mask=None):
+def nansum(
+    values: np.ndarray,
+    axis: Optional[int] = None,
+    skipna: bool = True,
+    min_count: int = 0,
+    mask: Optional[np.ndarray] = None,
+) -> Dtype:
     """
     Sum the elements along an axis ignoring NaNs
 
@@ -629,7 +652,7 @@ def _get_counts_nanvar(
     mask: Optional[np.ndarray],
     axis: Optional[int],
     ddof: int,
-    dtype=float,
+    dtype: Dtype = float,
 ) -> Tuple[Union[int, np.ndarray], Union[int, np.ndarray]]:
     """ Get the count of non-null values along an axis, accounting
     for degrees of freedom.
@@ -776,7 +799,13 @@ def nanvar(values, axis=None, skipna=True, ddof=1, mask=None):
 
 
 @disallow("M8", "m8")
-def nansem(values, axis=None, skipna=True, ddof=1, mask=None):
+def nansem(
+    values: np.ndarray,
+    axis: Optional[int] = None,
+    skipna: bool = True,
+    ddof: int = 1,
+    mask: Optional[np.ndarray] = None,
+) -> float:
     """
     Compute the standard error in the mean along given axis while ignoring NaNs
 
@@ -819,9 +848,14 @@ def nansem(values, axis=None, skipna=True, ddof=1, mask=None):
     return np.sqrt(var) / np.sqrt(count)
 
 
-def _nanminmax(meth, fill_value_typ):
+def _nanminmax(meth: str, fill_value_typ: str) -> Callable:
     @bottleneck_switch(name="nan" + meth)
-    def reduction(values, axis=None, skipna=True, mask=None):
+    def reduction(
+        values: np.ndarray,
+        axis: Optional[int] = None,
+        skipna: bool = True,
+        mask: Optional[np.ndarray] = None,
+    ) -> np.ndarray:
 
         values, mask, dtype, dtype_max, fill_value = _get_values(
             values, skipna, fill_value_typ=fill_value_typ, mask=mask
@@ -847,7 +881,12 @@ def reduction(values, axis=None, skipna=True, mask=None):
 
 
 @disallow("O")
-def nanargmax(values, axis=None, skipna=True, mask=None):
+def nanargmax(
+    values: np.ndarray,
+    axis: Optional[int] = None,
+    skipna: bool = True,
+    mask: Optional[np.ndarray] = None,
+) -> int:
     """
     Parameters
     ----------
@@ -878,7 +917,12 @@ def nanargmax(values, axis=None, skipna=True, mask=None):
 
 
 @disallow("O")
-def nanargmin(values, axis=None, skipna=True, mask=None):
+def nanargmin(
+    values: np.ndarray,
+    axis: Optional[int] = None,
+    skipna: bool = True,
+    mask: Optional[np.ndarray] = None,
+) -> int:
     """
     Parameters
     ----------
@@ -909,7 +953,12 @@ def nanargmin(values, axis=None, skipna=True, mask=None):
 
 
 @disallow("M8", "m8")
-def nanskew(values, axis=None, skipna=True, mask=None):
+def nanskew(
+    values: np.ndarray,
+    axis: Optional[int] = None,
+    skipna: bool = True,
+    mask: Optional[np.ndarray] = None,
+) -> float:
     """ Compute the sample skewness.
 
     The statistic computed here is the adjusted Fisher-Pearson standardized
@@ -987,7 +1036,12 @@ def nanskew(values, axis=None, skipna=True, mask=None):
 
 
 @disallow("M8", "m8")
-def nankurt(values, axis=None, skipna=True, mask=None):
+def nankurt(
+    values: np.ndarray,
+    axis: Optional[int] = None,
+    skipna: bool = True,
+    mask: Optional[np.ndarray] = None,
+) -> float:
     """
     Compute the sample excess kurtosis
 
@@ -1075,7 +1129,13 @@ def nankurt(values, axis=None, skipna=True, mask=None):
 
 
 @disallow("M8", "m8")
-def nanprod(values, axis=None, skipna=True, min_count=0, mask=None):
+def nanprod(
+    values: np.ndarray,
+    axis: Optional[int] = None,
+    skipna: bool = True,
+    min_count: int = 0,
+    mask: Optional[np.ndarray] = None,
+) -> Dtype:
     """
     Parameters
     ----------
@@ -1138,7 +1198,7 @@ def _get_counts(
     values_shape: Tuple[int],
     mask: Optional[np.ndarray],
     axis: Optional[int],
-    dtype=float,
+    dtype: Dtype = float,
 ) -> Union[int, np.ndarray]:
     """ Get the count of non-null values along an axis
 
@@ -1218,7 +1278,12 @@ def _zero_out_fperr(arg):
 
 
 @disallow("M8", "m8")
-def nancorr(a, b, method="pearson", min_periods=None):
+def nancorr(
+    a: np.ndarray,
+    b: np.ndarray,
+    method: str = "pearson",
+    min_periods: Optional[int] = None,
+):
     """
     a, b: ndarrays
     """
@@ -1240,7 +1305,7 @@ def nancorr(a, b, method="pearson", min_periods=None):
     return f(a, b)
 
 
-def get_corr_func(method):
+def get_corr_func(method: str):
     if method in ["kendall", "spearman"]:
         from scipy.stats import kendalltau, spearmanr
     elif callable(method):
@@ -1262,7 +1327,7 @@ def _spearman(a, b):
 
 
 @disallow("M8", "m8")
-def nancov(a, b, min_periods=None):
+def nancov(a: np.ndarray, b: np.ndarray, min_periods: Optional[int] = None):
     if len(a) != len(b):
         raise AssertionError("Operands to nancov must have same size")
 
@@ -1308,7 +1373,7 @@ def _ensure_numeric(x):
 # NA-friendly array comparisons
 
 
-def make_nancomp(op):
+def make_nancomp(op) -> Callable:
     def f(x, y):
         xmask = isna(x)
         ymask = isna(y)
@@ -1335,7 +1400,9 @@ def f(x, y):
 nanne = make_nancomp(operator.ne)
 
 
-def _nanpercentile_1d(values, mask, q, na_value, interpolation):
+def _nanpercentile_1d(
+    values: np.ndarray, mask: np.ndarray, q, na_value: Scalar, interpolation: str
+) -> Union[Scalar, np.ndarray]:
     """
     Wrapper for np.percentile that skips missing values, specialized to
     1-dimensional case.
@@ -1366,7 +1433,15 @@ def _nanpercentile_1d(values, mask, q, na_value, interpolation):
     return np.percentile(values, q, interpolation=interpolation)
 
 
-def nanpercentile(values, q, axis, na_value, mask, ndim, interpolation):
+def nanpercentile(
+    values: np.ndarray,
+    q,
+    axis: int,
+    na_value,
+    mask: np.ndarray,
+    ndim: int,
+    interpolation: str,
+):
     """
     Wrapper for np.percentile that skips missing values.