diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index f8913c5cc..c8e2cdff7 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -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 diff --git a/tests/test_frame.py b/tests/test_frame.py index 699f64877..99b323a67 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -3,6 +3,7 @@ from collections import defaultdict import csv import datetime +from enum import Enum import io import itertools from pathlib import Path @@ -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]})