From 1d649332ebbed677e714380881a0555bc6681a6a Mon Sep 17 00:00:00 2001 From: seljaks <33955366+seljaks@users.noreply.github.com> Date: Fri, 2 Dec 2022 16:05:12 +0100 Subject: [PATCH] ENH/TST: expand copy-on-write to assign GH49473 --- pandas/core/frame.py | 2 +- pandas/tests/copy_view/test_methods.py | 18 ++++++++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/pandas/core/frame.py b/pandas/core/frame.py index 0144aefedaa5f..5358fdb0b4dbd 100644 --- a/pandas/core/frame.py +++ b/pandas/core/frame.py @@ -4887,7 +4887,7 @@ def assign(self, **kwargs) -> DataFrame: Portland 17.0 62.6 290.15 Berkeley 25.0 77.0 298.15 """ - data = self.copy() + data = self.copy(deep=None) for k, v in kwargs.items(): data[k] = com.apply_if_callable(v, data) diff --git a/pandas/tests/copy_view/test_methods.py b/pandas/tests/copy_view/test_methods.py index bf65f153b10dd..6707f1411cbc7 100644 --- a/pandas/tests/copy_view/test_methods.py +++ b/pandas/tests/copy_view/test_methods.py @@ -280,3 +280,21 @@ def test_head_tail(method, using_copy_on_write): # without CoW enabled, head and tail return views. Mutating df2 also mutates df. df2.iloc[0, 0] = 1 tm.assert_frame_equal(df, df_orig) + + +def test_assign(using_copy_on_write): + df = DataFrame({"a": [1, 2, 3]}) + df_orig = df.copy() + df2 = df.assign() + df2._mgr._verify_integrity() + + if using_copy_on_write: + assert np.shares_memory(get_array(df2, "a"), get_array(df, "a")) + else: + assert not np.shares_memory(get_array(df2, "a"), get_array(df, "a")) + + # modify df2 to trigger CoW for that block + df2.iloc[0, 0] = 0 + if using_copy_on_write: + assert not np.shares_memory(get_array(df2, "a"), get_array(df, "a")) + tm.assert_frame_equal(df, df_orig)