Skip to content

Commit

Permalink
add Hashable type to __get_item__ #592 (#596)
Browse files Browse the repository at this point in the history
* add Hashable type to __get_item__ #592

* add assert_type to test_types_getitem_with_hashable

* code style: formatting

* remove tuple[Hashable, ...] from get_item overload1

* added more assert_type to assure other overloads are ok
  • Loading branch information
anilbey authored Mar 29, 2023
1 parent fbc5f20 commit e3d9a1f
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 1 deletion.
2 changes: 1 addition & 1 deletion pandas-stubs/core/frame.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -516,7 +516,7 @@ class DataFrame(NDFrame, OpsMixin):
def T(self) -> DataFrame: ...
def __getattr__(self, name: str) -> Series: ...
@overload
def __getitem__(self, idx: Scalar | tuple[Hashable, ...]) -> Series: ...
def __getitem__(self, idx: Scalar | Hashable) -> Series: ...
@overload
def __getitem__(self, rows: slice) -> DataFrame: ...
@overload
Expand Down
21 changes: 21 additions & 0 deletions tests/test_frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
from collections import defaultdict
import csv
import datetime
from enum import Enum
import io
import itertools
from pathlib import Path
Expand Down Expand Up @@ -160,6 +161,26 @@ def test_types_getitem() -> None:
df[i]


def test_types_getitem_with_hashable() -> None:
# Testing getitem support for hashable types that are not scalar
# Due to the bug in https://github.com/pandas-dev/pandas-stubs/issues/592
class MyEnum(Enum):
FIRST = "tayyar"
SECOND = "haydar"

df = pd.DataFrame(
data=[[12.2, 10], [8.8, 15]], columns=[MyEnum.FIRST, MyEnum.SECOND]
)
check(assert_type(df[MyEnum.FIRST], pd.Series), pd.Series)
check(assert_type(df[1:], pd.DataFrame), pd.DataFrame)
check(assert_type(df[:2], pd.DataFrame), pd.DataFrame)

df2 = pd.DataFrame(data=[[12.2, 10], [8.8, 15]], columns=[3, 4])
check(assert_type(df2[3], pd.Series), pd.Series)
check(assert_type(df2[[3]], pd.DataFrame), pd.DataFrame)
check(assert_type(df2[[3, 4]], pd.DataFrame), pd.DataFrame)


def test_slice_setitem() -> None:
# Due to the bug in pandas 1.2.3(https://github.com/pandas-dev/pandas/issues/40440), this is in separate test case
df = pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4], 5: [6, 7]})
Expand Down

0 comments on commit e3d9a1f

Please sign in to comment.