BUG: Series.loc raising KeyError for Iterator indexer in case of setitem (pandas-dev#39623)

Author: phofl

doc/source/whatsnew/v1.3.0.rst

@@ -691,6 +691,8 @@ Indexing
 - Bug in :meth:`DataFrame.__setitem__` raising ``ValueError`` with empty :class:`DataFrame` and specified columns for string indexer and non empty :class:`DataFrame` to set (:issue:`38831`)
 - Bug in :meth:`DataFrame.loc.__setitem__` raising ValueError when expanding unique column for :class:`DataFrame` with duplicate columns (:issue:`38521`)
 - Bug in :meth:`DataFrame.iloc.__setitem__` and :meth:`DataFrame.loc.__setitem__` with mixed dtypes when setting with a dictionary value (:issue:`38335`)
+- Bug in :meth:`Series.loc.__setitem__` and :meth:`DataFrame.loc.__setitem__` raising ``KeyError`` for boolean Iterator indexer (:issue:`39614`)
+- Bug in :meth:`Series.iloc` and :meth:`DataFrame.iloc` raising ``KeyError`` for Iterator indexer (:issue:`39614`)
 - Bug in :meth:`DataFrame.__setitem__` not raising ``ValueError`` when right hand side is a :class:`DataFrame` with wrong number of columns (:issue:`38604`)
 - Bug in :meth:`Series.__setitem__` raising ``ValueError`` when setting a :class:`Series` with a scalar indexer (:issue:`38303`)
 - Bug in :meth:`DataFrame.loc` dropping levels of :class:`MultiIndex` when :class:`DataFrame` used as input has only one row (:issue:`10521`)

pandas/core/indexing.py

@@ -703,6 +703,7 @@ def _ensure_listlike_indexer(self, key, axis=None, value=None):
 
     def __setitem__(self, key, value):
         if isinstance(key, tuple):
+            key = tuple(list(x) if is_iterator(x) else x for x in key)
             key = tuple(com.apply_if_callable(x, self.obj) for x in key)
         else:
             key = com.apply_if_callable(key, self.obj)
@@ -911,6 +912,7 @@ def _convert_to_indexer(self, key, axis: int, is_setter: bool = False):
 
     def __getitem__(self, key):
         if type(key) is tuple:
+            key = tuple(list(x) if is_iterator(x) else x for x in key)
             key = tuple(com.apply_if_callable(x, self.obj) for x in key)
             if self._is_scalar_access(key):
                 with suppress(KeyError, IndexError, AttributeError):
@@ -1242,6 +1244,9 @@ def _convert_to_indexer(self, key, axis: int, is_setter: bool = False):
 
         elif is_list_like_indexer(key):
 
+            if is_iterator(key):
+                key = list(key)
+
             if com.is_bool_indexer(key):
                 key = check_bool_indexer(labels, key)
                 (inds,) = key.nonzero()
@@ -1530,6 +1535,9 @@ def _getitem_axis(self, key, axis: int):
         if isinstance(key, slice):
             return self._get_slice_axis(key, axis=axis)
 
+        if is_iterator(key):
+            key = list(key)
+
         if isinstance(key, list):
             key = np.asarray(key)
 
@@ -1571,6 +1579,8 @@ def _convert_to_indexer(self, key, axis: int, is_setter: bool = False):
 
     def _get_setitem_indexer(self, key):
         # GH#32257 Fall through to let numpy do validation
+        if is_iterator(key):
+            return list(key)
         return key
 
     # -------------------------------------------------------------------

pandas/tests/frame/indexing/test_getitem.py

@@ -144,6 +144,21 @@ def test_getitem_listlike(self, idx_type, levels, float_frame):
         with pytest.raises(KeyError, match="not in index"):
             frame[idx]
 
+    def test_getitem_iloc_generator(self):
+        # GH#39614
+        df = DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
+        indexer = (x for x in [1, 2])
+        result = df.iloc[indexer]
+        expected = DataFrame({"a": [2, 3], "b": [5, 6]}, index=[1, 2])
+        tm.assert_frame_equal(result, expected)
+
+    def test_getitem_iloc_two_dimensional_generator(self):
+        df = DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
+        indexer = (x for x in [1, 2])
+        result = df.iloc[indexer, 1]
+        expected = Series([5, 6], name="b", index=[1, 2])
+        tm.assert_series_equal(result, expected)
+
 
 class TestGetitemCallable:
     def test_getitem_callable(self, float_frame):

pandas/tests/frame/indexing/test_setitem.py

@@ -610,6 +610,21 @@ def test_setitem_list_of_tuples(self, float_frame):
         expected = Series(tuples, index=float_frame.index, name="tuples")
         tm.assert_series_equal(result, expected)
 
+    def test_setitem_iloc_generator(self):
+        # GH#39614
+        df = DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
+        indexer = (x for x in [1, 2])
+        df.iloc[indexer] = 1
+        expected = DataFrame({"a": [1, 1, 1], "b": [4, 1, 1]})
+        tm.assert_frame_equal(df, expected)
+
+    def test_setitem_iloc_two_dimensional_generator(self):
+        df = DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
+        indexer = (x for x in [1, 2])
+        df.iloc[indexer, 1] = 1
+        expected = DataFrame({"a": [1, 2, 3], "b": [4, 1, 1]})
+        tm.assert_frame_equal(df, expected)
+
 
 class TestSetitemTZAwareValues:
     @pytest.fixture

pandas/tests/series/indexing/test_setitem.py

@@ -556,6 +556,9 @@ def test_int_key(self, obj, key, expected, val, indexer_sli, is_inplace):
         indkey = np.array(ilkey)
         self.check_indexer(obj, indkey, expected, val, indexer_sli, is_inplace)
 
+        genkey = (x for x in [key])
+        self.check_indexer(obj, genkey, expected, val, indexer_sli, is_inplace)
+
     def test_slice_key(self, obj, key, expected, val, indexer_sli, is_inplace):
         if not isinstance(key, slice):
             return
@@ -570,6 +573,9 @@ def test_slice_key(self, obj, key, expected, val, indexer_sli, is_inplace):
         indkey = np.array(ilkey)
         self.check_indexer(obj, indkey, expected, val, indexer_sli, is_inplace)
 
+        genkey = (x for x in indkey)
+        self.check_indexer(obj, genkey, expected, val, indexer_sli, is_inplace)
+
     def test_mask_key(self, obj, key, expected, val, indexer_sli):
         # setitem with boolean mask
         mask = np.zeros(obj.shape, dtype=bool)