Skip to content

887 select_dtypes stubs fixing #900

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Apr 2, 2024
Merged

887 select_dtypes stubs fixing #900

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
96ec420
adding tests for select_dtypes
Apr 2, 2024
eb6ffa6
adding tests for select_dtypes, resolves #887
Apr 2, 2024
ec64188
better handling of exclusions in arguments
Apr 2, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 31 additions & 3 deletions pandas-stubs/core/frame.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,10 @@ from pandas.core.window.rolling import (
Rolling,
Window,
)
from typing_extensions import Self
from typing_extensions import (
Self,
TypeAlias,
)
import xarray as xr

from pandas._libs.lib import NoDefault
Expand Down Expand Up @@ -608,10 +611,35 @@ 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]
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think if you added the following here as the first 2 overloads, you could catch the case of the string arguments not being allowed:

@overload
def select_dtypes(self, include: StrDtypeArg, exclude: AstypeArgExtList | None = ...) -> Never: ...
@overload
def select_dtypes(self, include: AstypeArgExtList | None = ..., exclude: StrDtypeArg) -> Never: ...

And if you added this, I think that you'd catch the case of 2 empty lists:

@overload
def select_dtypes(self, include: list[Never], exclude: list[Never]) -> Never: ...

You import Never from typing_extensions

There's a possibility that either mypy and/or pyright will complain about overlapping overloads, but you can add #type: ignore and # pyright: ignore as needed in the stubs to ignore that complaint.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That is implemented. Never is actually stronger than I expected since pylance read the stubs and visually show that the program will stop there.

@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,
Expand Down
52 changes: 52 additions & 0 deletions tests/test_frame.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3154,6 +3154,58 @@ 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:
# not able to check with typing that inputs lists are empty
# check(
# assert_type(df.select_dtypes([], []), pd.DataFrame), pd.DataFrame
# ) # ValueError

# ValueError :
check(assert_type(df.select_dtypes(), pd.DataFrame), pd.DataFrame) # type: ignore[assert-type, call-overload] # pyright: ignore[reportAssertTypeFailure, reportCallIssue]

# any kind of string dtype is not allowed but strings dtypes are included in AstypeArg...
# check(
# assert_type(df.select_dtypes(str), pd.DataFrame), pd.DataFrame
# ) # TypeError


def test_to_json_mode() -> None:
df = pd.DataFrame(
[["a", "b"], ["c", "d"]],
Expand Down
Loading