DataFrame__getitem__: accepts generator (#714)

685

Author: twoertwein

pandas-stubs/core/frame.pyi

@@ -1,5 +1,6 @@
 from collections.abc import (
     Callable,
+    Generator,
     Hashable,
     Iterable,
     Iterator,
@@ -550,14 +551,15 @@ class DataFrame(NDFrame, OpsMixin):
     def T(self) -> DataFrame: ...
     def __getattr__(self, name: str) -> Series: ...
     @overload
-    def __getitem__(
+    def __getitem__(  # type: ignore[misc]
         self,
         key: Series[_bool]
         | DataFrame
         | Index
         | np_ndarray_str
         | np_ndarray_bool
-        | list[_ScalarOrTupleT],
+        | list[_ScalarOrTupleT]
+        | Generator[_ScalarOrTupleT, None, None],
     ) -> DataFrame: ...
     @overload
     def __getitem__(self, key: slice) -> DataFrame: ...

tests/test_frame.py

@@ -2715,3 +2715,8 @@ def test_groupby_fillna_inplace() -> None:
     check(assert_type(groupby.fillna(0, inplace=False), pd.DataFrame), pd.DataFrame)
     if TYPE_CHECKING_INVALID_USAGE:
         groupby.fillna(0, inplace=True)  # type: ignore[arg-type] # pyright: ignore[reportGeneralTypeIssues]
+
+
+def test_getitem_generator() -> None:
+    # GH 685
+    check(assert_type(DF[(f"col{i+1}" for i in range(2))], pd.DataFrame), pd.DataFrame)