Skip to content

Commit

Permalink
Use the new TypeIs feature to update the notna/notnull/isna/isnull ty…
Browse files Browse the repository at this point in the history
…pe guards (#972)

* GH 927 Allow 'integer' in select_dtypes include argument

* TypeGuard clean up in dtypes/missing.pyi

* PR Feedback
  • Loading branch information
loicdiridollou authored Aug 7, 2024
1 parent 1b4fd6f commit 1896191
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 34 deletions.
6 changes: 3 additions & 3 deletions pandas-stubs/core/dtypes/missing.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ from pandas import (
Index,
Series,
)
from typing_extensions import TypeGuard
from typing_extensions import TypeIs

from pandas._libs.missing import NAType
from pandas._libs.tslibs import NaTType
Expand All @@ -32,7 +32,7 @@ def isna(obj: Index[Any] | list[Any] | ArrayLike) -> npt.NDArray[np.bool_]: ...
@overload
def isna(
obj: Scalar | NaTType | NAType | None,
) -> TypeGuard[NaTType | NAType | None]: ...
) -> TypeIs[NaTType | NAType | None]: ...

isnull = isna

Expand All @@ -43,6 +43,6 @@ def notna(obj: Series[Any]) -> Series[bool]: ...
@overload
def notna(obj: Index[Any] | list[Any] | ArrayLike) -> npt.NDArray[np.bool_]: ...
@overload
def notna(obj: ScalarT | NaTType | NAType | None) -> TypeGuard[ScalarT]: ...
def notna(obj: ScalarT | NaTType | NAType | None) -> TypeIs[ScalarT]: ...

notnull = notna
40 changes: 9 additions & 31 deletions tests/test_pandas.py
Original file line number Diff line number Diff line change
Expand Up @@ -366,60 +366,38 @@ def test_isna() -> None:
assert check(assert_type(pd.isna(np_nat), bool), bool)
assert not check(assert_type(pd.notna(np_nat), 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
# Check TypeIs type narrowing functionality
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)
check(assert_type(nullable1, str), 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])
assert_type(nullable1, Union[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)
check(assert_type(nullable2, int), 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))
check(assert_type(nullable2, 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
check(assert_type(nullable2, None), type(None))

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)
check(assert_type(nullable3, bool), bool)
if pd.isna(nullable3):
# assert_type(nullable3, Union[NAType, None]) # TODO: Desired result (see comments above)
assert_type(nullable3, Union[NaTType, NAType, None])
assert_type(nullable3, Union[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
assert_type(nullable3, Union[NAType, None])


# GH 55
Expand Down

0 comments on commit 1896191

Please sign in to comment.