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]

pandas/core/internals/array_manager.py

@@ -21,6 +21,7 @@
 from pandas._typing import (
     ArrayLike,
     DtypeObj,
+    npt,
 )
 from pandas.util._validators import validate_bool_kwarg
 
@@ -568,7 +569,7 @@ def reindex_indexer(
     def _reindex_indexer(
         self: T,
         new_axis,
-        indexer,
+        indexer: npt.NDArray[np.intp] | None,
         axis: int,
         fill_value=None,
         allow_dups: bool = False,
@@ -579,7 +580,7 @@ def _reindex_indexer(
         Parameters
         ----------
         new_axis : Index
-        indexer : ndarray of int64 or None
+        indexer : ndarray[intp] or None
         axis : int
         fill_value : object, default None
         allow_dups : bool, default False

pandas/core/internals/blocks.py

@@ -87,7 +87,6 @@
     replace_regex,
     should_use_regex,
 )
-from pandas.core.array_algos.take import take_nd
 from pandas.core.array_algos.transforms import shift
 from pandas.core.arrays import (
     Categorical,
@@ -736,7 +735,7 @@ def _replace_coerce(
         self,
         to_replace,
         value,
-        mask: np.ndarray,
+        mask: npt.NDArray[np.bool_],
         inplace: bool = True,
         regex: bool = False,
     ) -> list[Block]:
@@ -827,19 +826,14 @@ def set_inplace(self, locs, values: ArrayLike) -> None:
 
     def take_nd(
         self,
-        indexer,
+        indexer: npt.NDArray[np.intp],
         axis: int,
         new_mgr_locs: BlockPlacement | None = None,
         fill_value=lib.no_default,
     ) -> Block:
         """
-        Take values according to indexer and return them as a block.bb
-
+        Take values according to indexer and return them as a block.
         """
-        # algos.take_nd dispatches for DatetimeTZBlock, CategoricalBlock
-        # so need to preserve types
-        # sparse is treated like an ndarray, but needs .get_values() shaping
-
         values = self.values
 
         if fill_value is lib.no_default:
@@ -848,6 +842,7 @@ def take_nd(
         else:
             allow_fill = True
 
+        # Note: algos.take_nd has upcast logic similar to coerce_to_target_dtype
         new_values = algos.take_nd(
             values, indexer, axis=axis, allow_fill=allow_fill, fill_value=fill_value
         )
@@ -1727,7 +1722,7 @@ def is_numeric(self):
 
     def take_nd(
         self,
-        indexer,
+        indexer: npt.NDArray[np.intp],
         axis: int = 0,
         new_mgr_locs: BlockPlacement | None = None,
         fill_value=lib.no_default,
@@ -2259,7 +2254,7 @@ def to_native_types(
     """convert to our native types format"""
     if isinstance(values, Categorical):
         # GH#40754 Convert categorical datetimes to datetime array
-        values = take_nd(
+        values = algos.take_nd(
             values.categories._values,
             ensure_platform_int(values._codes),
             fill_value=na_rep,

pandas/core/internals/managers.py

@@ -597,7 +597,7 @@ def consolidate(self: T) -> T:
     def reindex_indexer(
         self: T,
         new_axis: Index,
-        indexer,
+        indexer: npt.NDArray[np.intp] | None,
         axis: int,
         fill_value=None,
         allow_dups: bool = False,
@@ -610,7 +610,7 @@ def reindex_indexer(
         Parameters
         ----------
         new_axis : Index
-        indexer : ndarray of int64 or None
+        indexer : ndarray[intp] or None
         axis : int
         fill_value : object, default None
         allow_dups : bool, default False

pandas/core/missing.py

@@ -43,7 +43,7 @@
     from pandas import Index
 
 
-def check_value_size(value, mask: np.ndarray, length: int):
+def check_value_size(value, mask: npt.NDArray[np.bool_], length: int):
     """
     Validate the size of the values passed to ExtensionArray.fillna.
     """