diff --git a/pandas-stubs/core/dtypes/missing.pyi b/pandas-stubs/core/dtypes/missing.pyi index 747f260bd..84bc8ff35 100644 --- a/pandas-stubs/core/dtypes/missing.pyi +++ b/pandas-stubs/core/dtypes/missing.pyi @@ -1,7 +1,4 @@ -from typing import ( - Literal, - overload, -) +from typing import overload import numpy as np from numpy import typing as npt @@ -10,12 +7,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 +27,9 @@ 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]: ... +def isna( + obj: Scalar | NaTType | NAType | None, +) -> TypeGuard[NaTType | NAType | None]: ... isnull = isna @@ -41,8 +40,6 @@ 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]: ... +def notna(obj: ScalarT | NaTType | NAType | None) -> TypeGuard[ScalarT]: ... notnull = notna diff --git a/tests/test_pandas.py b/tests/test_pandas.py index d5f73991b..456f94947 100644 --- a/tests/test_pandas.py +++ b/tests/test_pandas.py @@ -5,7 +5,6 @@ from typing import ( TYPE_CHECKING, Any, - Literal, Union, ) @@ -17,6 +16,8 @@ import pytest from typing_extensions import assert_type +from pandas._libs.missing import NAType +from pandas._libs.tslibs import NaTType from pandas._typing import Scalar from tests import ( @@ -246,17 +247,72 @@ def test_isna() -> None: idx2 = pd.Index([1, 2]) check(assert_type(pd.notna(idx2), npt.NDArray[np.bool_]), np.ndarray, np.bool_) - assert check(assert_type(pd.isna(pd.NA), Literal[True]), bool) - assert not check(assert_type(pd.notna(pd.NA), Literal[False]), bool) - - assert check(assert_type(pd.isna(pd.NaT), Literal[True]), bool) - assert not check(assert_type(pd.notna(pd.NaT), Literal[False]), bool) - - assert check(assert_type(pd.isna(None), Literal[True]), bool) - assert not check(assert_type(pd.notna(None), Literal[False]), bool) - - check(assert_type(pd.isna(2.5), bool), bool) - check(assert_type(pd.notna(2.5), bool), bool) + assert check(assert_type(pd.isna(pd.NA), bool), bool) + assert not check(assert_type(pd.notna(pd.NA), bool), bool) + + assert check(assert_type(pd.isna(pd.NaT), bool), bool) + assert not check(assert_type(pd.notna(pd.NaT), bool), bool) + + assert check(assert_type(pd.isna(None), bool), bool) + assert not check(assert_type(pd.notna(None), bool), bool) + + assert not check(assert_type(pd.isna(2.5), bool), bool) + assert check(assert_type(pd.notna(2.5), bool), bool) + + # Check TypeGuard type narrowing functionality + # TODO: Due to limitations in TypeGuard spec, the true annotations are not always viable + # and as a result the type narrowing does not always work as it intuitively should + # There is a proposal being floated for a StrictTypeGuard that will have more rigid narrowing semantics + # In the test cases below, a commented out assertion will be included to document the optimal test result + nullable1: str | None | NAType | NaTType = random.choice( + ["value", None, pd.NA, pd.NaT] + ) + if pd.notna(nullable1): + check(assert_type(nullable1, str), str) + if not pd.isna(nullable1): + # check(assert_type(nullable1, str), str) # TODO: Desired result (see comments above) + check(assert_type(nullable1, Union[str, NaTType, NAType, None]), str) + if pd.isna(nullable1): + assert_type(nullable1, Union[NaTType, NAType, None]) + if not pd.notna(nullable1): + # assert_type(nullable1, Union[NaTType, NAType, None]) # TODO: Desired result (see comments above) + assert_type(nullable1, Union[str, NaTType, NAType, None]) + + nullable2: int | None = random.choice([2, None]) + if pd.notna(nullable2): + check(assert_type(nullable2, int), int) + if not pd.isna(nullable2): + # check(assert_type(nullable2, int), int) # TODO: Desired result (see comments above) + check(assert_type(nullable2, Union[int, None]), int) + if pd.isna(nullable2): + # check(assert_type(nullable2, None), type(None)) # TODO: Desired result (see comments above) + check(assert_type(nullable2, Union[NaTType, NAType, None]), type(None)) + if not pd.notna(nullable2): + # check(assert_type(nullable2, None), type(None)) # TODO: Desired result (see comments above) + # TODO: MyPy and Pyright produce conflicting results: + # assert_type(nullable2, Union[int, None]) # MyPy result + # assert_type( + # nullable2, Union[int, NaTType, NAType, None] + # ) # Pyright result + pass + + nullable3: bool | None | NAType = random.choice([True, None, pd.NA]) + if pd.notna(nullable3): + check(assert_type(nullable3, bool), bool) + if not pd.isna(nullable3): + # check(assert_type(nullable3, bool), bool) # TODO: Desired result (see comments above) + check(assert_type(nullable3, Union[bool, NAType, None]), bool) + if pd.isna(nullable3): + # assert_type(nullable3, Union[NAType, None]) # TODO: Desired result (see comments above) + assert_type(nullable3, Union[NaTType, NAType, None]) + if not pd.notna(nullable3): + # assert_type(nullable3, Union[NAType, None]) # TODO: Desired result (see comments above) + # TODO: MyPy and Pyright produce conflicting results: + # assert_type(nullable3, Union[bool, NAType, None]) # Mypy result + # assert_type( + # nullable3, Union[bool, NaTType, NAType, None] + # ) # Pyright result + pass # GH 55