[backport 2.3.x] BUG(CoW): also raise for chained assignment for .at / .iat (#62074) (#62115)

Author: jorisvandenbossche

pandas/core/indexing.py

@@ -2575,6 +2575,22 @@ def __getitem__(self, key):
         return super().__getitem__(key)
 
     def __setitem__(self, key, value) -> None:
+        if not PYPY and using_copy_on_write():
+            if sys.getrefcount(self.obj) <= 2:
+                warnings.warn(
+                    _chained_assignment_msg, ChainedAssignmentError, stacklevel=2
+                )
+        elif not PYPY and not using_copy_on_write():
+            ctr = sys.getrefcount(self.obj)
+            ref_count = 2
+            if not warn_copy_on_write() and _check_cacher(self.obj):
+                # see https://github.com/pandas-dev/pandas/pull/56060#discussion_r1399245221
+                ref_count += 1
+            if ctr <= ref_count:
+                warnings.warn(
+                    _chained_assignment_warning_msg, FutureWarning, stacklevel=2
+                )
+
         if self.ndim == 2 and not self._axes_are_unique:
             # GH#33041 fall back to .loc
             if not isinstance(key, tuple) or not all(is_scalar(x) for x in key):
@@ -2599,6 +2615,25 @@ def _convert_key(self, key):
                 raise ValueError("iAt based indexing can only have integer indexers")
         return key
 
+    def __setitem__(self, key, value) -> None:
+        if not PYPY and using_copy_on_write():
+            if sys.getrefcount(self.obj) <= 2:
+                warnings.warn(
+                    _chained_assignment_msg, ChainedAssignmentError, stacklevel=2
+                )
+        elif not PYPY and not using_copy_on_write():
+            ctr = sys.getrefcount(self.obj)
+            ref_count = 2
+            if not warn_copy_on_write() and _check_cacher(self.obj):
+                # see https://github.com/pandas-dev/pandas/pull/56060#discussion_r1399245221
+                ref_count += 1
+            if ctr <= ref_count:
+                warnings.warn(
+                    _chained_assignment_warning_msg, FutureWarning, stacklevel=2
+                )
+
+        return super().__setitem__(key, value)
+
 
 def _tuplify(ndim: int, loc: Hashable) -> tuple[Hashable | slice, ...]:
     """

pandas/tests/copy_view/test_chained_assignment_deprecation.py

@@ -172,3 +172,83 @@ def test_frame_setitem(indexer, using_copy_on_write):
     with option_context("chained_assignment", "warn"):
         with tm.raises_chained_assignment_error(extra_warnings=extra_warnings):
             df[0:3][indexer] = 10
+
+
+@pytest.mark.parametrize(
+    "indexer", [0, [0, 1], slice(0, 2), np.array([True, False, True])]
+)
+def test_series_iloc_setitem(indexer):
+    df = DataFrame({"a": [1, 2, 3], "b": 1})
+
+    with option_context("chained_assignment", "warn"):
+        with tm.raises_chained_assignment_error():
+            df["a"].iloc[indexer] = 0
+
+
+@pytest.mark.parametrize(
+    "indexer", [0, [0, 1], slice(0, 2), np.array([True, False, True])]
+)
+def test_frame_iloc_setitem(indexer, using_copy_on_write):
+    df = DataFrame({"a": [1, 2, 3, 4, 5], "b": 1})
+
+    extra_warnings = () if using_copy_on_write else (SettingWithCopyWarning,)
+
+    with option_context("chained_assignment", "warn"):
+        with tm.raises_chained_assignment_error(extra_warnings=extra_warnings):
+            df[0:3].iloc[indexer] = 10
+
+
+@pytest.mark.parametrize(
+    "indexer", [0, [0, 1], slice(0, 2), np.array([True, False, True])]
+)
+def test_series_loc_setitem(indexer):
+    df = DataFrame({"a": [1, 2, 3], "b": 1})
+
+    with option_context("chained_assignment", "warn"):
+        with tm.raises_chained_assignment_error():
+            df["a"].loc[indexer] = 0
+
+
+@pytest.mark.parametrize(
+    "indexer", [0, [0, 1], (0, "a"), slice(0, 2), np.array([True, False, True])]
+)
+def test_frame_loc_setitem(indexer, using_copy_on_write):
+    df = DataFrame({"a": [1, 2, 3, 4, 5], "b": 1})
+
+    extra_warnings = () if using_copy_on_write else (SettingWithCopyWarning,)
+
+    with option_context("chained_assignment", "warn"):
+        with tm.raises_chained_assignment_error(extra_warnings=extra_warnings):
+            df[0:3].loc[indexer] = 10
+
+
+def test_series_at_setitem():
+    df = DataFrame({"a": [1, 2, 3], "b": 1})
+
+    with option_context("chained_assignment", "warn"):
+        with tm.raises_chained_assignment_error():
+            df["a"].at[0] = 0
+
+
+def test_frame_at_setitem():
+    df = DataFrame({"a": [1, 2, 3, 4, 5], "b": 1})
+
+    with option_context("chained_assignment", "warn"):
+        with tm.raises_chained_assignment_error():
+            df[0:3].at[0, "a"] = 10
+
+
+def test_series_iat_setitem():
+    df = DataFrame({"a": [1, 2, 3], "b": 1})
+
+    with option_context("chained_assignment", "warn"):
+        with tm.raises_chained_assignment_error():
+            df["a"].iat[0] = 0
+
+
+def test_frame_iat_setitem():
+    df = DataFrame({"a": [1, 2, 3, 4, 5], "b": 1})
+
+    with option_context("chained_assignment", "warn"):
+        with tm.raises_chained_assignment_error():
+            df[0:3].iat[0, 0] = 10