From c2bb19d8cfbfa2e375c39c960cc7315f43717965 Mon Sep 17 00:00:00 2001 From: Loic Diridollou Date: Sun, 4 Aug 2024 15:24:55 -0700 Subject: [PATCH 1/3] GH 927 Allow 'integer' in select_dtypes include argument --- pandas-stubs/core/frame.pyi | 1 + tests/test_frame.py | 1 + 2 files changed, 2 insertions(+) diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 20879b000..8607343c7 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -624,6 +624,7 @@ class DataFrame(NDFrame, OpsMixin): "number", "datetime64", "datetime", + "integer", "timedelta", "timedelta64", "datetimetz", diff --git a/tests/test_frame.py b/tests/test_frame.py index 4b68f391c..32f23e240 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -3226,6 +3226,7 @@ def test_convert_dtypes_dtype_backend() -> None: def test_select_dtypes() -> None: df = pd.DataFrame({"a": [1, 2] * 3, "b": [True, False] * 3, "c": [1.0, 2.0] * 3}) check(assert_type(df.select_dtypes("number"), pd.DataFrame), pd.DataFrame) + check(assert_type(df.select_dtypes("integer"), pd.DataFrame), pd.DataFrame) check(assert_type(df.select_dtypes(np.number), pd.DataFrame), pd.DataFrame) check(assert_type(df.select_dtypes(object), pd.DataFrame), pd.DataFrame) check(assert_type(df.select_dtypes(include="bool"), pd.DataFrame), pd.DataFrame) From df7dad9275d780cba1c2c1d413c42c8c7e775d32 Mon Sep 17 00:00:00 2001 From: Loic Diridollou Date: Tue, 6 Aug 2024 16:32:24 -0700 Subject: [PATCH 2/3] TypeGuard clean up in dtypes/missing.pyi --- pandas-stubs/core/dtypes/missing.pyi | 6 +++--- tests/test_pandas.py | 27 ++++++++------------------- 2 files changed, 11 insertions(+), 22 deletions(-) diff --git a/pandas-stubs/core/dtypes/missing.pyi b/pandas-stubs/core/dtypes/missing.pyi index 22bf8edd8..11319304d 100644 --- a/pandas-stubs/core/dtypes/missing.pyi +++ b/pandas-stubs/core/dtypes/missing.pyi @@ -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 @@ -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 @@ -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 diff --git a/tests/test_pandas.py b/tests/test_pandas.py index e72538673..04439323e 100644 --- a/tests/test_pandas.py +++ b/tests/test_pandas.py @@ -366,52 +366,41 @@ 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) + check(assert_type(nullable2, None), type(None)) # 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) + 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: From edfdef14a91eea520fbeddfecef83a4b48432e54 Mon Sep 17 00:00:00 2001 From: Loic Diridollou Date: Tue, 6 Aug 2024 16:46:03 -0700 Subject: [PATCH 3/3] PR Feedback --- tests/test_pandas.py | 13 +------------ 1 file changed, 1 insertion(+), 12 deletions(-) diff --git a/tests/test_pandas.py b/tests/test_pandas.py index 04439323e..3cd36af71 100644 --- a/tests/test_pandas.py +++ b/tests/test_pandas.py @@ -388,11 +388,6 @@ def test_isna() -> None: check(assert_type(nullable2, None), type(None)) if not pd.notna(nullable2): check(assert_type(nullable2, None), type(None)) - # 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 nullable3: bool | None | NAType = random.choice([True, None, pd.NA]) if pd.notna(nullable3): @@ -402,13 +397,7 @@ def test_isna() -> None: if pd.isna(nullable3): 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