add Hashable type to __get_item__ #592

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

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]})