remove types from Index.__iter__()` (#532)

* remove types from Index.__iter__()`

* add some tests

Author: Dr-Irv

pandas-stubs/core/indexes/base.pyi

@@ -33,7 +33,6 @@ from pandas._typing import (
     DtypeObj,
     FillnaOptions,
     HashableT,
-    IndexIterScalar,
     IndexT,
     Label,
     Level,
@@ -223,7 +222,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[IndexIterScalar | tuple[Hashable, ...]]: ...
+    def __iter__(self) -> Iterator: ...
     def __ne__(self, other: object) -> np_ndarray_bool: ...  # type: ignore[override]
     def __le__(self, other: Index | Scalar) -> np_ndarray_bool: ...  # type: ignore[override]
     def __ge__(self, other: Index | Scalar) -> np_ndarray_bool: ...  # type: ignore[override]

tests/test_frame.py

@@ -1813,7 +1813,6 @@ def test_frame_index_numpy() -> None:
 
 
 def test_frame_stack() -> None:
-
     multicol2 = pd.MultiIndex.from_tuples([("weight", "kg"), ("height", "m")])
     df_multi_level_cols2 = pd.DataFrame(
         [[1.0, 2.0], [3.0, 4.0]], index=["cat", "dog"], columns=multicol2
@@ -2403,3 +2402,13 @@ def get_NDArray(df: pd.DataFrame, key: npt.NDArray[np.uint64]) -> pd.DataFrame:
 
     a: npt.NDArray[np.uint64] = np.array([10, 30], dtype="uint64")
     check(assert_type(get_NDArray(df, a), pd.DataFrame), pd.DataFrame)
+
+
+def test_in_columns() -> None:
+    # GH 532 (PR)
+    df = pd.DataFrame(np.random.random((3, 4)), columns=["cat", "dog", "rat", "pig"])
+    cols = [c for c in df.columns if "at" in c]
+    check(assert_type(cols, list), list, str)
+    check(assert_type(df.loc[:, cols], pd.DataFrame), pd.DataFrame)
+    check(assert_type(df[cols], pd.DataFrame), pd.DataFrame)
+    check(assert_type(df.groupby(by=cols).sum(), pd.DataFrame), pd.DataFrame)

tests/test_indexes.py

@@ -3,8 +3,6 @@
 import datetime as dt
 from typing import (
     TYPE_CHECKING,
-    Hashable,
-    List,
     Tuple,
     Union,
 )
@@ -29,7 +27,6 @@
 if TYPE_CHECKING:
     from pandas.core.indexes.numeric import NumericIndex
 
-    from pandas._typing import IndexIterScalar
 else:
     if not PD_LTE_15:
         from pandas import Index as NumericIndex
@@ -108,7 +105,7 @@ def test_column_sequence() -> None:
     df = pd.DataFrame([1, 2, 3])
     col_list = list(df.columns)
     check(
-        assert_type(col_list, List[Union["IndexIterScalar", Tuple[Hashable, ...]]]),
+        assert_type(col_list, list),
         list,
         int,
     )
@@ -684,14 +681,14 @@ def test_sorted_and_list() -> None:
     check(
         assert_type(
             sorted(i1),
-            List[Union["IndexIterScalar", Tuple[Hashable, ...]]],
+            list,
         ),
         list,
     )
     check(
         assert_type(
             list(i1),
-            List[Union["IndexIterScalar", Tuple[Hashable, ...]]],
+            list,
         ),
         list,
     )