Skip to content

Simplify DataFrame.__getitem__ #771

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 4 commits into from
Sep 1, 2023
Merged

Simplify DataFrame.__getitem__ #771

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
319ca58
simplify DataFrame.__getitem__
twoertwein Sep 1, 2023
2181f29
restrict to Iterable with Hashable/Scalar elements
twoertwein Sep 1, 2023
596cd82
unrelated: widen type of DataFrame.__iter__, can be any column type
twoertwein Sep 1, 2023
6e78f8c
Scalar | Hashable -> Hashable
twoertwein Sep 1, 2023
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
22 changes: 4 additions & 18 deletions pandas-stubs/core/frame.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
from collections.abc import (
Callable,
Generator,
Hashable,
Iterable,
Iterator,
Expand All @@ -14,7 +13,6 @@ from typing import (
Any,
ClassVar,
Literal,
TypeVar,
overload,
)

Expand Down Expand Up @@ -119,8 +117,6 @@ from pandas._typing import (
ValidationOptions,
WriteBuffer,
XMLParsers,
np_ndarray_bool,
np_ndarray_str,
npt,
num,
)
Expand All @@ -130,7 +126,6 @@ from pandas.plotting import PlotAccessor

_str = str
_bool = bool
_ScalarOrTupleT = TypeVar("_ScalarOrTupleT", bound=Scalar | tuple[Hashable, ...])

class _iLocIndexerFrame(_iLocIndexer):
@overload
Expand Down Expand Up @@ -553,20 +548,11 @@ class DataFrame(NDFrame, OpsMixin):
def T(self) -> DataFrame: ...
def __getattr__(self, name: str) -> Series: ...
@overload
def __getitem__( # type: ignore[misc]
self,
key: Series
| DataFrame
| Index
| np_ndarray_str
| np_ndarray_bool
| list[_ScalarOrTupleT]
| Generator[_ScalarOrTupleT, None, None],
) -> DataFrame: ...
def __getitem__(self, key: Scalar | tuple[Hashable, ...]) -> Series: ... # type: ignore[misc]
@overload
def __getitem__(self, key: slice) -> DataFrame: ...
def __getitem__(self, key: Iterable[Scalar | Hashable] | slice) -> DataFrame: ...
Copy link
Collaborator

Choose a reason for hiding this comment

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

Could this be simplified to just Iterable[Hashable], since anything that is Scalar is Hashable ? (Same for __iter__())

Copy link
Member Author

Choose a reason for hiding this comment

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

Yes, that is possible - changed it!

@overload
def __getitem__(self, key: Scalar | Hashable) -> Series: ...
def __getitem__(self, key: Hashable) -> Series: ...
def isetitem(
self, loc: int | Sequence[int], value: Scalar | ArrayLike | list[Any]
) -> None: ...
Expand Down Expand Up @@ -1477,7 +1463,7 @@ class DataFrame(NDFrame, OpsMixin):
Name: _str
#
# dunder methods
def __iter__(self) -> Iterator[float | _str]: ...
def __iter__(self) -> Iterator[Scalar | Hashable]: ...
# properties
@property
def at(self): ... # Not sure what to do with this yet; look at source
Expand Down
7 changes: 7 additions & 0 deletions tests/test_frame.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2811,3 +2811,10 @@ def test_groupby_fillna_inplace() -> None:
def test_getitem_generator() -> None:
# GH 685
check(assert_type(DF[(f"col{i+1}" for i in range(2))], pd.DataFrame), pd.DataFrame)


def test_getitem_dict_keys() -> None:
# GH 770
some_columns = {"a": [1], "b": [2]}
df = pd.DataFrame.from_dict(some_columns)
check(assert_type(df[some_columns.keys()], pd.DataFrame), pd.DataFrame)