TYP: assorted (#46121)

Author: jbrockmendel

pandas/core/algorithms.py

@@ -534,7 +534,7 @@ def factorize_array(
     na_sentinel: int = -1,
     size_hint: int | None = None,
     na_value=None,
-    mask: np.ndarray | None = None,
+    mask: npt.NDArray[np.bool_] | None = None,
 ) -> tuple[npt.NDArray[np.intp], np.ndarray]:
     """
     Factorize a numpy array to codes and uniques.

pandas/core/array_algos/masked_reductions.py

@@ -2,27 +2,26 @@
 masked_reductions.py is for reduction algorithms using a mask-based approach
 for missing values.
 """
+from __future__ import annotations
 
-from typing import (
-    Callable,
-    Optional,
-)
+from typing import Callable
 
 import numpy as np
 
 from pandas._libs import missing as libmissing
+from pandas._typing import npt
 
 from pandas.core.nanops import check_below_min_count
 
 
 def _sumprod(
     func: Callable,
     values: np.ndarray,
-    mask: np.ndarray,
+    mask: npt.NDArray[np.bool_],
     *,
     skipna: bool = True,
     min_count: int = 0,
-    axis: Optional[int] = None,
+    axis: int | None = None,
 ):
     """
     Sum or product for 1D masked array.
@@ -33,7 +32,7 @@ def _sumprod(
     values : np.ndarray
         Numpy array with the values (can be of any dtype that support the
         operation).
-    mask : np.ndarray
+    mask : np.ndarray[bool]
         Boolean numpy array (True values indicate missing values).
     skipna : bool, default True
         Whether to skip NA.
@@ -58,11 +57,11 @@ def _sumprod(
 
 def sum(
     values: np.ndarray,
-    mask: np.ndarray,
+    mask: npt.NDArray[np.bool_],
     *,
     skipna: bool = True,
     min_count: int = 0,
-    axis: Optional[int] = None,
+    axis: int | None = None,
 ):
     return _sumprod(
         np.sum, values=values, mask=mask, skipna=skipna, min_count=min_count, axis=axis
@@ -71,11 +70,11 @@ def sum(
 
 def prod(
     values: np.ndarray,
-    mask: np.ndarray,
+    mask: npt.NDArray[np.bool_],
     *,
     skipna: bool = True,
     min_count: int = 0,
-    axis: Optional[int] = None,
+    axis: int | None = None,
 ):
     return _sumprod(
         np.prod, values=values, mask=mask, skipna=skipna, min_count=min_count, axis=axis
@@ -85,10 +84,10 @@ def prod(
 def _minmax(
     func: Callable,
     values: np.ndarray,
-    mask: np.ndarray,
+    mask: npt.NDArray[np.bool_],
     *,
     skipna: bool = True,
-    axis: Optional[int] = None,
+    axis: int | None = None,
 ):
     """
     Reduction for 1D masked array.
@@ -99,7 +98,7 @@ def _minmax(
     values : np.ndarray
         Numpy array with the values (can be of any dtype that support the
         operation).
-    mask : np.ndarray
+    mask : np.ndarray[bool]
         Boolean numpy array (True values indicate missing values).
     skipna : bool, default True
         Whether to skip NA.
@@ -122,26 +121,26 @@ def _minmax(
 
 def min(
     values: np.ndarray,
-    mask: np.ndarray,
+    mask: npt.NDArray[np.bool_],
     *,
     skipna: bool = True,
-    axis: Optional[int] = None,
+    axis: int | None = None,
 ):
     return _minmax(np.min, values=values, mask=mask, skipna=skipna, axis=axis)
 
 
 def max(
     values: np.ndarray,
-    mask: np.ndarray,
+    mask: npt.NDArray[np.bool_],
     *,
     skipna: bool = True,
-    axis: Optional[int] = None,
+    axis: int | None = None,
 ):
     return _minmax(np.max, values=values, mask=mask, skipna=skipna, axis=axis)
 
 
 # TODO: axis kwarg
-def mean(values: np.ndarray, mask: np.ndarray, skipna: bool = True):
+def mean(values: np.ndarray, mask: npt.NDArray[np.bool_], skipna: bool = True):
     if not values.size or mask.all():
         return libmissing.NA
     _sum = _sumprod(np.sum, values=values, mask=mask, skipna=skipna)

pandas/core/array_algos/replace.py

@@ -15,6 +15,7 @@
 from pandas._typing import (
     ArrayLike,
     Scalar,
+    npt,
 )
 
 from pandas.core.dtypes.common import (
@@ -42,7 +43,7 @@ def should_use_regex(regex: bool, to_replace: Any) -> bool:
 
 
 def compare_or_regex_search(
-    a: ArrayLike, b: Scalar | Pattern, regex: bool, mask: np.ndarray
+    a: ArrayLike, b: Scalar | Pattern, regex: bool, mask: npt.NDArray[np.bool_]
 ) -> ArrayLike | bool:
     """
     Compare two array-like inputs of the same shape or two scalar values
@@ -116,7 +117,9 @@ def _check_comparison_types(
     return result
 
 
-def replace_regex(values: ArrayLike, rx: re.Pattern, value, mask: np.ndarray | None):
+def replace_regex(
+    values: ArrayLike, rx: re.Pattern, value, mask: npt.NDArray[np.bool_] | None
+):
     """
     Parameters
     ----------

pandas/core/arrays/_mixins.py

@@ -367,7 +367,7 @@ def _putmask(self, mask: npt.NDArray[np.bool_], value) -> None:
         np.putmask(self._ndarray, mask, value)
 
     def _where(
-        self: NDArrayBackedExtensionArrayT, mask: np.ndarray, value
+        self: NDArrayBackedExtensionArrayT, mask: npt.NDArray[np.bool_], value
     ) -> NDArrayBackedExtensionArrayT:
         """
         Analogue to np.where(mask, self, value)

pandas/core/arrays/masked.py

@@ -99,13 +99,15 @@ class BaseMaskedArray(OpsMixin, ExtensionArray):
     _internal_fill_value: Scalar
     # our underlying data and mask are each ndarrays
     _data: np.ndarray
-    _mask: np.ndarray
+    _mask: npt.NDArray[np.bool_]
 
     # Fill values used for any/all
     _truthy_value = Scalar  # bool(_truthy_value) = True
     _falsey_value = Scalar  # bool(_falsey_value) = False
 
-    def __init__(self, values: np.ndarray, mask: np.ndarray, copy: bool = False):
+    def __init__(
+        self, values: np.ndarray, mask: npt.NDArray[np.bool_], copy: bool = False
+    ):
         # values is supposed to already be validated in the subclass
         if not (isinstance(mask, np.ndarray) and mask.dtype == np.bool_):
             raise TypeError(

pandas/core/arrays/numeric.py

@@ -17,6 +17,7 @@
 from pandas._typing import (
     Dtype,
     DtypeObj,
+    npt,
 )
 from pandas.errors import AbstractMethodError
 from pandas.util._decorators import cache_readonly
@@ -219,7 +220,9 @@ class NumericArray(BaseMaskedArray):
 
     _dtype_cls: type[NumericDtype]
 
-    def __init__(self, values: np.ndarray, mask: np.ndarray, copy: bool = False):
+    def __init__(
+        self, values: np.ndarray, mask: npt.NDArray[np.bool_], copy: bool = False
+    ):
         checker = self._dtype_cls._checker
         if not (isinstance(values, np.ndarray) and checker(values.dtype)):
             descr = (

pandas/core/groupby/groupby.py

@@ -3607,13 +3607,13 @@ def tail(self, n=5):
         return self._mask_selected_obj(mask)
 
     @final
-    def _mask_selected_obj(self, mask: np.ndarray) -> NDFrameT:
+    def _mask_selected_obj(self, mask: npt.NDArray[np.bool_]) -> NDFrameT:
         """
         Return _selected_obj with mask applied to the correct axis.
 
         Parameters
         ----------
-        mask : np.ndarray
+        mask : np.ndarray[bool]
             Boolean mask to apply.
 
         Returns

pandas/core/groupby/ops.py

@@ -437,8 +437,8 @@ def _cython_op_ndim_compat(
         min_count: int,
         ngroups: int,
         comp_ids: np.ndarray,
-        mask: np.ndarray | None = None,
-        result_mask: np.ndarray | None = None,
+        mask: npt.NDArray[np.bool_] | None = None,
+        result_mask: npt.NDArray[np.bool_] | None = None,
         **kwargs,
     ) -> np.ndarray:
         if values.ndim == 1:
@@ -481,8 +481,8 @@ def _call_cython_op(
         min_count: int,
         ngroups: int,
         comp_ids: np.ndarray,
-        mask: np.ndarray | None,
-        result_mask: np.ndarray | None,
+        mask: npt.NDArray[np.bool_] | None,
+        result_mask: npt.NDArray[np.bool_] | None,
         **kwargs,
     ) -> np.ndarray:  # np.ndarray[ndim=2]
         orig_values = values

pandas/core/indexes/base.py

@@ -5548,7 +5548,9 @@ def asof(self, label):
 
         return self[loc]
 
-    def asof_locs(self, where: Index, mask: np.ndarray) -> npt.NDArray[np.intp]:
+    def asof_locs(
+        self, where: Index, mask: npt.NDArray[np.bool_]
+    ) -> npt.NDArray[np.intp]:
         """
         Return the locations (indices) of labels in the index.
 

pandas/core/indexes/period.py

@@ -23,6 +23,7 @@
 from pandas._typing import (
     Dtype,
     DtypeObj,
+    npt,
 )
 from pandas.util._decorators import doc
 from pandas.util._exceptions import find_stack_level
@@ -327,7 +328,7 @@ def _is_comparable_dtype(self, dtype: DtypeObj) -> bool:
     # ------------------------------------------------------------------------
     # Index Methods
 
-    def asof_locs(self, where: Index, mask: np.ndarray) -> np.ndarray:
+    def asof_locs(self, where: Index, mask: npt.NDArray[np.bool_]) -> np.ndarray:
         """
         where : array of timestamps
         mask : np.ndarray[bool]