Turn isna() and notna() into TypeGuards

Author: gandhis1

pandas-stubs/core/dtypes/missing.pyi

@@ -10,12 +10,14 @@ from pandas import (
     Index,
     Series,
 )
+from typing_extensions import TypeGuard
 
 from pandas._libs.missing import NAType
 from pandas._libs.tslibs import NaTType
 from pandas._typing import (
     ArrayLike,
     Scalar,
+    ScalarT,
 )
 
 isposinf_scalar = ...
@@ -28,9 +30,17 @@ def isna(obj: Series) -> Series[bool]: ...
 @overload
 def isna(obj: Index | list | ArrayLike) -> npt.NDArray[np.bool_]: ...
 @overload
-def isna(obj: Scalar) -> bool: ...
-@overload
 def isna(obj: NaTType | NAType | None) -> Literal[True]: ...
+@overload
+def isna(
+    obj: Scalar | NaTType | NAType | None,
+) -> TypeGuard[NaTType | NAType | None]: ...
+@overload
+def isna(obj: Scalar | NaTType) -> TypeGuard[NaTType]: ...
+@overload
+def isna(obj: Scalar | NAType) -> TypeGuard[NAType]: ...
+@overload
+def isna(obj: Scalar | None) -> TypeGuard[None]: ...
 
 isnull = isna
 
@@ -41,8 +51,8 @@ def notna(obj: Series) -> Series[bool]: ...
 @overload
 def notna(obj: Index | list | ArrayLike) -> npt.NDArray[np.bool_]: ...
 @overload
-def notna(obj: Scalar) -> bool: ...
-@overload
 def notna(obj: NaTType | NAType | None) -> Literal[False]: ...
+@overload
+def notna(obj: ScalarT | NaTType | NAType | None) -> TypeGuard[ScalarT]: ...
 
 notnull = notna

tests/test_pandas.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import random
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -13,6 +14,9 @@
 from pandas.api.extensions import ExtensionArray
 from typing_extensions import assert_type
 
+from pandas._libs.missing import NAType
+from pandas._libs.tslibs import NaTType
+
 from tests import check
 
 
@@ -142,6 +146,25 @@ def test_isna() -> None:
     check(assert_type(pd.isna(2.5), bool), bool)
     check(assert_type(pd.notna(2.5), bool), bool)
 
+    # Check type guard functionality
+    nullable1 = random.choice(["value", None, pd.NA, pd.NaT])
+    if pd.notna(nullable1):
+        check(assert_type(nullable1, str), str)
+    if pd.isna(nullable1):
+        assert_type(nullable1, Union[NaTType, NAType, None])
+
+    nullable2 = random.choice([2, None])
+    if pd.notna(nullable2):
+        check(assert_type(nullable2, int), int)
+    if pd.isna(nullable2):
+        assert_type(nullable2, None)
+
+    nullable3 = random.choice([2, None, pd.NA])
+    if pd.notna(nullable3):
+        check(assert_type(nullable3, int), int)
+    if pd.isna(nullable3):
+        assert_type(nullable3, Union[NAType, None])
+
 
 # GH 55
 def test_read_xml() -> None: