Fix return type of __iter__() and tuple-handling for __getitem__()

Author: gandhis1

pandas-stubs/core/frame.pyi

@@ -11,6 +11,7 @@ from typing import (
     Mapping,
     Pattern,
     Sequence,
+    TypeVar,
     overload,
 )
 
@@ -115,6 +116,7 @@ from pandas.plotting import PlotAccessor
 
 _str = str
 _bool = bool
+_ScalarOrTupleT = TypeVar("_ScalarOrTupleT", bound=Scalar | tuple[Hashable, ...])
 
 class _iLocIndexerFrame(_iLocIndexer):
     @overload
@@ -491,21 +493,18 @@ class DataFrame(NDFrame, OpsMixin):
     def T(self) -> DataFrame: ...
     def __getattr__(self, name: str) -> Series: ...
     @overload
-    def __getitem__(self, idx: Scalar) -> Series: ...
+    def __getitem__(self, idx: Scalar | tuple[Hashable, ...]) -> Series: ...
     @overload
     def __getitem__(self, rows: slice) -> DataFrame: ...
     @overload
     def __getitem__(
         self,
-        idx: tuple
-        | Series[_bool]
+        idx: Series[_bool]
         | DataFrame
-        | list[_str]
-        | list[ScalarT]
         | Index
         | np_ndarray_str
         | np_ndarray_bool
-        | Sequence[tuple[Scalar, ...]],
+        | list[_ScalarOrTupleT],
     ) -> DataFrame: ...
     def isetitem(
         self, loc: int | Sequence[int], value: Scalar | ArrayLike | list[Any]

pandas-stubs/core/indexes/base.pyi

@@ -220,7 +220,7 @@ class Index(IndexOpsMixin, PandasObject):
     def shape(self) -> tuple[int, ...]: ...
     # Extra methods from old stubs
     def __eq__(self, other: object) -> np_ndarray_bool: ...  # type: ignore[override]
-    def __iter__(self) -> Iterator: ...
+    def __iter__(self) -> Iterator[Scalar | tuple[Hashable, ...]]: ...
     def __ne__(self, other: object) -> np_ndarray_bool: ...  # type: ignore[override]
     def __le__(self, other: Index | Scalar) -> np_ndarray_bool: ...
     def __ge__(self, other: Index | Scalar) -> np_ndarray_bool: ...

tests/test_frame.py

@@ -1465,7 +1465,9 @@ def test_getmultiindex_columns() -> None:
     res3: pd.DataFrame = df[
         [(i, s) for i in [1] for s in df.columns.get_level_values(1)]
     ]
-    ndf: pd.DataFrame = df[[df.columns[0]]]
+    res4: pd.DataFrame = df[[df.columns[0]]]
+    check(assert_type(df[df.columns[0]], pd.Series), pd.Series)
+    check(assert_type(df[li[0]], pd.Series), pd.Series)
 
 
 def test_frame_getitem_isin() -> None:

tests/test_indexes.py

@@ -1,14 +1,21 @@
 from __future__ import annotations
 
 import datetime as dt
-from typing import Union
+from typing import (
+    Hashable,
+    List,
+    Tuple,
+    Union,
+)
 
 import numpy as np
 from numpy import typing as npt
 import pandas as pd
 from pandas.core.indexes.numeric import NumericIndex
 from typing_extensions import assert_type
 
+from pandas._typing import Scalar
+
 from tests import (
     check,
     pytest_warns_bounded,
@@ -67,6 +74,7 @@ def test_column_getitem() -> None:
     df = pd.DataFrame([[1, 2, 3]], columns=["a", "b", "c"])
 
     column = df.columns[0]
+    check(assert_type(column, Scalar), str)
     check(assert_type(df[column], pd.Series), pd.Series, int)
 
 
@@ -81,6 +89,12 @@ def test_column_contains() -> None:
     length = len(df.columns[df.columns.str.contains("A|B")])
 
 
+def test_column_sequence() -> None:
+    df = pd.DataFrame([1, 2, 3])
+    col_list = list(df.columns)
+    check(assert_type(col_list, List[Union[Scalar, Tuple[Hashable, ...]]]), list, int)
+
+
 def test_difference_none() -> None:
     # https://github.com/pandas-dev/pandas-stubs/issues/17
     ind = pd.Index([1, 2, 3])