TYP: pd.isna (#46222)

Author: twoertwein

pandas/core/arrays/boolean.py

@@ -1,7 +1,10 @@
 from __future__ import annotations
 
 import numbers
-from typing import TYPE_CHECKING
+from typing import (
+    TYPE_CHECKING,
+    cast,
+)
 
 import numpy as np
 
@@ -31,6 +34,8 @@
 if TYPE_CHECKING:
     import pyarrow
 
+    from pandas._typing import npt
+
 
 @register_extension_dtype
 class BooleanDtype(BaseMaskedDtype):
@@ -188,7 +193,9 @@ def coerce_to_array(
         if inferred_dtype not in ("boolean", "empty") + integer_like:
             raise TypeError("Need to pass bool-like values")
 
-        mask_values = isna(values_object)
+        # mypy does not narrow the type of mask_values to npt.NDArray[np.bool_]
+        # within this branch, it assumes it can also be None
+        mask_values = cast("npt.NDArray[np.bool_]", isna(values_object))
         values = np.zeros(len(values), dtype=bool)
         values[~mask_values] = values_object[~mask_values].astype(bool)
 

pandas/core/base.py

@@ -769,7 +769,9 @@ def hasnans(self) -> bool:
 
         Enables various performance speedups.
         """
-        return bool(isna(self).any())
+        # error: Item "bool" of "Union[bool, ndarray[Any, dtype[bool_]], NDFrame]"
+        # has no attribute "any"
+        return bool(isna(self).any())  # type: ignore[union-attr]
 
     def isna(self):
         return isna(self._values)

pandas/core/dtypes/missing.py

@@ -5,6 +5,10 @@
 
 from decimal import Decimal
 from functools import partial
+from typing import (
+    TYPE_CHECKING,
+    overload,
+)
 
 import numpy as np
 
@@ -16,11 +20,6 @@
     NaT,
     iNaT,
 )
-from pandas._typing import (
-    ArrayLike,
-    DtypeObj,
-    npt,
-)
 
 from pandas.core.dtypes.common import (
     DT64NS_DTYPE,
@@ -54,6 +53,19 @@
 )
 from pandas.core.dtypes.inference import is_list_like
 
+if TYPE_CHECKING:
+    from pandas._typing import (
+        ArrayLike,
+        DtypeObj,
+        NDFrame,
+        NDFrameT,
+        Scalar,
+        npt,
+    )
+
+    from pandas.core.indexes.base import Index
+
+
 isposinf_scalar = libmissing.isposinf_scalar
 isneginf_scalar = libmissing.isneginf_scalar
 
@@ -63,7 +75,35 @@
 _dtype_str = np.dtype(str)
 
 
-def isna(obj):
+@overload
+def isna(obj: Scalar) -> bool:
+    ...
+
+
+@overload
+def isna(
+    obj: ArrayLike | Index | list,
+) -> npt.NDArray[np.bool_]:
+    ...
+
+
+@overload
+def isna(obj: NDFrameT) -> NDFrameT:
+    ...
+
+
+# handle unions
+@overload
+def isna(obj: NDFrameT | ArrayLike | Index | list) -> NDFrameT | npt.NDArray[np.bool_]:
+    ...
+
+
+@overload
+def isna(obj: object) -> bool | npt.NDArray[np.bool_] | NDFrame:
+    ...
+
+
+def isna(obj: object) -> bool | npt.NDArray[np.bool_] | NDFrame:
     """
     Detect missing values for an array-like object.
 
@@ -284,7 +324,35 @@ def _isna_string_dtype(values: np.ndarray, inf_as_na: bool) -> npt.NDArray[np.bo
     return result
 
 
-def notna(obj):
+@overload
+def notna(obj: Scalar) -> bool:
+    ...
+
+
+@overload
+def notna(
+    obj: ArrayLike | Index | list,
+) -> npt.NDArray[np.bool_]:
+    ...
+
+
+@overload
+def notna(obj: NDFrameT) -> NDFrameT:
+    ...
+
+
+# handle unions
+@overload
+def notna(obj: NDFrameT | ArrayLike | Index | list) -> NDFrameT | npt.NDArray[np.bool_]:
+    ...
+
+
+@overload
+def notna(obj: object) -> bool | npt.NDArray[np.bool_] | NDFrame:
+    ...
+
+
+def notna(obj: object) -> bool | npt.NDArray[np.bool_] | NDFrame:
     """
     Detect non-missing values for an array-like object.
 
@@ -362,7 +430,7 @@ def notna(obj):
     Name: 1, dtype: bool
     """
     res = isna(obj)
-    if is_scalar(res):
+    if isinstance(res, bool):
         return not res
     return ~res
 

pandas/core/window/ewm.py

@@ -3,7 +3,10 @@
 import datetime
 from functools import partial
 from textwrap import dedent
-from typing import TYPE_CHECKING
+from typing import (
+    TYPE_CHECKING,
+    cast,
+)
 import warnings
 
 import numpy as np
@@ -380,12 +383,11 @@ def __init__(
                     FutureWarning,
                     stacklevel=find_stack_level(),
                 )
-                self.times = self._selected_obj[self.times]
+                # self.times cannot be str anymore
+                self.times = cast("Series", self._selected_obj[self.times])
             if not is_datetime64_ns_dtype(self.times):
                 raise ValueError("times must be datetime64[ns] dtype.")
-            # error: Argument 1 to "len" has incompatible type "Union[str, ndarray,
-            # NDFrameT, None]"; expected "Sized"
-            if len(self.times) != len(obj):  # type: ignore[arg-type]
+            if len(self.times) != len(obj):
                 raise ValueError("times must be the same length as the object.")
             if not isinstance(self.halflife, (str, datetime.timedelta)):
                 raise ValueError(

pandas/io/parsers/base_parser.py

@@ -855,9 +855,11 @@ def _check_data_length(
         data: list of array-likes containing the data column-wise.
         """
         if not self.index_col and len(columns) != len(data) and columns:
-            if len(columns) == len(data) - 1 and np.all(
-                (is_object_dtype(data[-1]) and data[-1] == "") | isna(data[-1])
-            ):
+            empty_str = is_object_dtype(data[-1]) and data[-1] == ""
+            # error: No overload variant of "__ror__" of "ndarray" matches
+            # argument type "ExtensionArray"
+            empty_str_or_na = empty_str | isna(data[-1])  # type: ignore[operator]
+            if len(columns) == len(data) - 1 and np.all(empty_str_or_na):
                 return
             warnings.warn(
                 "Length of header or names does not match length of data. This leads "