From 91e251c3234333efadd0b467020bfcd300d9b6d8 Mon Sep 17 00:00:00 2001
From: Patrick Hoefler <61934744+phofl@users.noreply.github.com>
Date: Fri, 8 Dec 2023 10:38:38 +0100
Subject: [PATCH] CoW: Avoid warning in case of expansion (#56391)

---
 pandas/core/internals/managers.py            | 17 ++++++++++++-----
 pandas/tests/copy_view/test_indexing.py      |  3 +--
 pandas/tests/frame/indexing/test_indexing.py |  4 +---
 3 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/pandas/core/internals/managers.py b/pandas/core/internals/managers.py
index a221d02b75bb2..cc88312d5b58f 100644
--- a/pandas/core/internals/managers.py
+++ b/pandas/core/internals/managers.py
@@ -1325,16 +1325,16 @@ def column_setitem(
         This is a method on the BlockManager level, to avoid creating an
         intermediate Series at the DataFrame level (`s = df[loc]; s[idx] = value`)
         """
+        needs_to_warn = False
         if warn_copy_on_write() and not self._has_no_reference(loc):
             if not isinstance(
                 self.blocks[self.blknos[loc]].values,
                 (ArrowExtensionArray, ArrowStringArray),
             ):
-                warnings.warn(
-                    COW_WARNING_GENERAL_MSG,
-                    FutureWarning,
-                    stacklevel=find_stack_level(),
-                )
+                # We might raise if we are in an expansion case, so defer
+                # warning till we actually updated
+                needs_to_warn = True
+
         elif using_copy_on_write() and not self._has_no_reference(loc):
             blkno = self.blknos[loc]
             # Split blocks to only copy the column we want to modify
@@ -1358,6 +1358,13 @@ def column_setitem(
             new_mgr = col_mgr.setitem((idx,), value)
             self.iset(loc, new_mgr._block.values, inplace=True)
 
+        if needs_to_warn:
+            warnings.warn(
+                COW_WARNING_GENERAL_MSG,
+                FutureWarning,
+                stacklevel=find_stack_level(),
+            )
+
     def insert(self, loc: int, item: Hashable, value: ArrayLike, refs=None) -> None:
         """
         Insert item at selected position.
diff --git a/pandas/tests/copy_view/test_indexing.py b/pandas/tests/copy_view/test_indexing.py
index 33b8ce218f029..6f3850ab64daa 100644
--- a/pandas/tests/copy_view/test_indexing.py
+++ b/pandas/tests/copy_view/test_indexing.py
@@ -1145,13 +1145,12 @@ def test_set_value_copy_only_necessary_column(
     view = df[:]
 
     if val == "a" and indexer[0] != slice(None):
-        # TODO(CoW-warn) assert the FutureWarning for CoW is also raised
         with tm.assert_produces_warning(
             FutureWarning, match="Setting an item of incompatible dtype is deprecated"
         ):
             indexer_func(df)[indexer] = val
     else:
-        with tm.assert_cow_warning(warn_copy_on_write):
+        with tm.assert_cow_warning(warn_copy_on_write and val == 100):
             indexer_func(df)[indexer] = val
 
     if using_copy_on_write:
diff --git a/pandas/tests/frame/indexing/test_indexing.py b/pandas/tests/frame/indexing/test_indexing.py
index dfb4a3092789a..40c6b8e180c5b 100644
--- a/pandas/tests/frame/indexing/test_indexing.py
+++ b/pandas/tests/frame/indexing/test_indexing.py
@@ -1397,9 +1397,7 @@ def test_iloc_setitem_enlarge_no_warning(self, warn_copy_on_write):
         df = DataFrame(columns=["a", "b"])
         expected = df.copy()
         view = df[:]
-        # TODO(CoW-warn) false positive: shouldn't warn in case of enlargement?
-        with tm.assert_produces_warning(FutureWarning if warn_copy_on_write else None):
-            df.iloc[:, 0] = np.array([1, 2], dtype=np.float64)
+        df.iloc[:, 0] = np.array([1, 2], dtype=np.float64)
         tm.assert_frame_equal(view, expected)
 
     def test_loc_internals_not_updated_correctly(self):