887 select_dtypes stubs fixing (pandas-dev#900)

* adding tests for select_dtypes

* adding tests for select_dtypes, resolves pandas-dev#887

* better handling of exclusions in arguments

---------

Co-authored-by: Laurent Mutricy <[email protected]>

Author: mutricyl

pandas-stubs/core/frame.pyi

@@ -50,7 +50,11 @@ from pandas.core.window.rolling import (
     Rolling,
     Window,
 )
-from typing_extensions import Self
+from typing_extensions import (
+    Never,
+    Self,
+    TypeAlias,
+)
 import xarray as xr
 
 from pandas._libs.lib import NoDefault
@@ -112,6 +116,7 @@ from pandas._typing import (
     SortKind,
     StataDateFormat,
     StorageOptions,
+    StrDtypeArg,
     StrLike,
     Suffixes,
     T as _T,
@@ -608,10 +613,47 @@ class DataFrame(NDFrame, OpsMixin):
         self, expr: _str, *, inplace: Literal[False] = ..., **kwargs
     ) -> DataFrame: ...
     def eval(self, expr: _str, *, inplace: _bool = ..., **kwargs): ...
+    AstypeArgExt: TypeAlias = (
+        AstypeArg
+        | Literal[
+            "number",
+            "datetime64",
+            "datetime",
+            "timedelta",
+            "timedelta64",
+            "datetimetz",
+            "datetime64[ns]",
+        ]
+    )
+    AstypeArgExtList: TypeAlias = AstypeArgExt | list[AstypeArgExt]
+    @overload
+    def select_dtypes(
+        self, include: StrDtypeArg, exclude: AstypeArgExtList | None = ...
+    ) -> Never: ...
+    @overload
+    def select_dtypes(
+        self, include: AstypeArgExtList | None, exclude: StrDtypeArg
+    ) -> Never: ...
+    @overload
+    def select_dtypes(self, exclude: StrDtypeArg) -> Never: ...
+    @overload
+    def select_dtypes(self, include: list[Never], exclude: list[Never]) -> Never: ...
+    @overload
+    def select_dtypes(
+        self,
+        include: AstypeArgExtList,
+        exclude: AstypeArgExtList | None = ...,
+    ) -> DataFrame: ...
+    @overload
+    def select_dtypes(
+        self,
+        include: AstypeArgExtList | None,
+        exclude: AstypeArgExtList,
+    ) -> DataFrame: ...
+    @overload
     def select_dtypes(
         self,
-        include: _str | list[_str] | None = ...,
-        exclude: _str | list[_str] | None = ...,
+        exclude: AstypeArgExtList,
     ) -> DataFrame: ...
     def insert(
         self,

tests/test_frame.py

@@ -38,6 +38,7 @@
 import pytest
 from typing_extensions import (
     TypeAlias,
+    assert_never,
     assert_type,
 )
 import xarray as xr
@@ -3154,6 +3155,53 @@ def test_convert_dtypes_dtype_backend() -> None:
     check(assert_type(dfn, pd.DataFrame), pd.DataFrame)
 
 
+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(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)
+    check(
+        assert_type(df.select_dtypes(include=["float64"], exclude=None), pd.DataFrame),
+        pd.DataFrame,
+    )
+    check(
+        assert_type(df.select_dtypes(exclude=["int64"], include=None), pd.DataFrame),
+        pd.DataFrame,
+    )
+    check(
+        assert_type(df.select_dtypes(exclude=["int64", object]), pd.DataFrame),
+        pd.DataFrame,
+    )
+    check(
+        assert_type(
+            df.select_dtypes(
+                exclude=[
+                    np.datetime64,
+                    "datetime64",
+                    "datetime",
+                    np.timedelta64,
+                    "timedelta",
+                    "timedelta64",
+                    "category",
+                    "datetimetz",
+                    "datetime64[ns]",
+                ]
+            ),
+            pd.DataFrame,
+        ),
+        pd.DataFrame,
+    )
+    if TYPE_CHECKING_INVALID_USAGE:
+        # include and exclude shall not be both empty
+        assert_never(df.select_dtypes([], []))
+        assert_never(df.select_dtypes())
+        # str like dtypes are not allowed
+        assert_never(df.select_dtypes(str))
+        assert_never(df.select_dtypes(exclude=str))
+        assert_never(df.select_dtypes(None, str))
+
+
 def test_to_json_mode() -> None:
     df = pd.DataFrame(
         [["a", "b"], ["c", "d"]],