API: detect and raise error for chained assignment under Copy-on-Write

Author: jorisvandenbossche

pandas/core/frame.py

@@ -16,6 +16,7 @@
 import functools
 from io import StringIO
 import itertools
+import sys
 from textwrap import dedent
 from typing import (
     TYPE_CHECKING,
@@ -95,7 +96,10 @@
     function as nv,
     np_percentile_argname,
 )
-from pandas.errors import InvalidIndexError
+from pandas.errors import (
+    ChainedAssignmentError,
+    InvalidIndexError,
+)
 from pandas.util._decorators import (
     Appender,
     Substitution,
@@ -3838,6 +3842,13 @@ def isetitem(self, loc, value) -> None:
         self._iset_item_mgr(loc, arraylike, inplace=False)
 
     def __setitem__(self, key, value):
+        if (
+            get_option("mode.copy_on_write")
+            and get_option("mode.data_manager") == "block"
+        ):
+            if sys.getrefcount(self) <= 3:
+                raise ChainedAssignmentError("Chained assignment doesn't work!!")
+
         key = com.apply_if_callable(key, self)
 
         # see if we can slice the rows

pandas/core/indexing.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from contextlib import suppress
+import sys
 from typing import (
     TYPE_CHECKING,
     Hashable,
@@ -12,6 +13,8 @@
 
 import numpy as np
 
+from pandas._config import get_option
+
 from pandas._libs.indexing import NDFrameIndexerBase
 from pandas._libs.lib import item_from_zerodim
 from pandas._typing import (
@@ -20,6 +23,7 @@
 )
 from pandas.errors import (
     AbstractMethodError,
+    ChainedAssignmentError,
     IndexingError,
     InvalidIndexError,
     LossySetitemError,
@@ -830,6 +834,14 @@ def _ensure_listlike_indexer(self, key, axis=None, value=None) -> None:
 
     @final
     def __setitem__(self, key, value) -> None:
+        if (
+            get_option("mode.copy_on_write")
+            and get_option("mode.data_manager") == "block"
+        ):
+            print("_LocationIndexer.__setitem__ refcount: ", sys.getrefcount(self.obj))
+            if sys.getrefcount(self.obj) <= 2:
+                raise ChainedAssignmentError("Chained assignment doesn't work!!")
+
         check_dict_or_set_indexers(key)
         if isinstance(key, tuple):
             key = tuple(list(x) if is_iterator(x) else x for x in key)

pandas/core/series.py

@@ -3,6 +3,7 @@
 """
 from __future__ import annotations
 
+import sys
 from textwrap import dedent
 from typing import (
     IO,
@@ -68,7 +69,10 @@
     npt,
 )
 from pandas.compat.numpy import function as nv
-from pandas.errors import InvalidIndexError
+from pandas.errors import (
+    ChainedAssignmentError,
+    InvalidIndexError,
+)
 from pandas.util._decorators import (
     Appender,
     Substitution,
@@ -1070,6 +1074,14 @@ def _get_value(self, label, takeable: bool = False):
             return self.iloc[loc]
 
     def __setitem__(self, key, value) -> None:
+        if (
+            get_option("mode.copy_on_write")
+            and get_option("mode.data_manager") == "block"
+        ):
+            print("Series.__getitem__ refcount: ", sys.getrefcount(self))
+            if sys.getrefcount(self) <= 3:
+                raise ChainedAssignmentError("Chained assignment doesn't work!!")
+
         check_dict_or_set_indexers(key)
         key = com.apply_if_callable(key, self)
         cacher_needs_updating = self._check_is_chained_assignment_possible()

pandas/errors/__init__.py

@@ -298,6 +298,28 @@ class SettingWithCopyError(ValueError):
     """
 
 
+class ChainedAssignmentError(ValueError):
+    """
+    Exception raised when trying to set on a copied slice from a ``DataFrame``.
+
+    The ``mode.chained_assignment`` needs to be set to set to 'raise.' This can
+    happen unintentionally when chained indexing.
+
+    For more information on eveluation order,
+    see :ref:`the user guide<indexing.evaluation_order>`.
+
+    For more information on view vs. copy,
+    see :ref:`the user guide<indexing.view_versus_copy>`.
+
+    Examples
+    --------
+    >>> pd.options.mode.chained_assignment = 'raise'
+    >>> df = pd.DataFrame({'A': [1, 1, 1, 2, 2]}, columns=['A'])
+    >>> df.loc[0:3]['A'] = 'a' # doctest: +SKIP
+    ... # SettingWithCopyError: A value is trying to be set on a copy of a...
+    """
+
+
 class SettingWithCopyWarning(Warning):
     """
     Warning raised when trying to set on a copied slice from a ``DataFrame``.

pandas/tests/frame/indexing/test_setitem.py

@@ -3,6 +3,7 @@
 import numpy as np
 import pytest
 
+from pandas.errors import ChainedAssignmentError
 import pandas.util._test_decorators as td
 
 from pandas.core.dtypes.base import _registry as ea_registry
@@ -1126,7 +1127,7 @@ def test_setitem_ea_boolean_mask(self):
 
 
 class TestDataFrameSetitemCopyViewSemantics:
-    def test_setitem_always_copy(self, float_frame):
+    def test_setitem_always_copy(self, float_frame, using_copy_on_write):
         assert "E" not in float_frame.columns
         s = float_frame["A"].copy()
         float_frame["E"] = s
@@ -1245,12 +1246,15 @@ def test_setitem_column_update_inplace(self, using_copy_on_write):
         df = DataFrame({col: np.zeros(len(labels)) for col in labels}, index=labels)
         values = df._mgr.blocks[0].values
 
-        for label in df.columns:
-            df[label][label] = 1
-
         if not using_copy_on_write:
+            for label in df.columns:
+                df[label][label] = 1
+
             # diagonal values all updated
             assert np.all(values[np.arange(10), np.arange(10)] == 1)
         else:
+            with pytest.raises(ChainedAssignmentError):
+                for label in df.columns:
+                    df[label][label] = 1
             # original dataframe not updated
             assert np.all(values[np.arange(10), np.arange(10)] == 0)

pandas/tests/frame/indexing/test_xs.py

@@ -3,7 +3,10 @@
 import numpy as np
 import pytest
 
-from pandas.errors import SettingWithCopyError
+from pandas.errors import (
+    ChainedAssignmentError,
+    SettingWithCopyError,
+)
 
 from pandas import (
     DataFrame,
@@ -124,7 +127,8 @@ def test_xs_view(self, using_array_manager, using_copy_on_write):
         df_orig = dm.copy()
 
         if using_copy_on_write:
-            dm.xs(2)[:] = 20
+            with pytest.raises(ChainedAssignmentError):
+                dm.xs(2)[:] = 20
             tm.assert_frame_equal(dm, df_orig)
         elif using_array_manager:
             # INFO(ArrayManager) with ArrayManager getting a row as a view is

pandas/tests/frame/methods/test_sort_values.py

@@ -345,8 +345,8 @@ def test_sort_values_frame_column_inplace_sort_exception(
             with pytest.raises(ValueError, match="This Series is a view"):
                 s.sort_values(inplace=True)
 
-        cp = s.copy()
-        cp.sort_values()  # it works!
+            cp = s.copy()
+            cp.sort_values()  # it works!
 
     def test_sort_values_nat_values_in_int_column(self):
 

pandas/tests/frame/test_block_internals.py

@@ -7,7 +7,10 @@
 import numpy as np
 import pytest
 
-from pandas.errors import PerformanceWarning
+from pandas.errors import (
+    ChainedAssignmentError,
+    PerformanceWarning,
+)
 import pandas.util._test_decorators as td
 
 import pandas as pd
@@ -340,7 +343,11 @@ def test_stale_cached_series_bug_473(self, using_copy_on_write):
             )
             repr(Y)
             Y["e"] = Y["e"].astype("object")
-            Y["g"]["c"] = np.NaN
+            if using_copy_on_write:
+                with pytest.raises(ChainedAssignmentError):
+                    Y["g"]["c"] = np.NaN
+            else:
+                Y["g"]["c"] = np.NaN
             repr(Y)
             result = Y.sum()  # noqa
             exp = Y["g"].sum()  # noqa

pandas/tests/indexing/multiindex/test_chaining_and_caching.py

@@ -1,7 +1,10 @@
 import numpy as np
 import pytest
 
-from pandas.errors import SettingWithCopyError
+from pandas.errors import (
+    ChainedAssignmentError,
+    SettingWithCopyError,
+)
 import pandas.util._test_decorators as td
 
 from pandas import (
@@ -50,11 +53,13 @@ def test_cache_updating(using_copy_on_write):
 
     # setting via chained assignment
     # but actually works, since everything is a view
-    df.loc[0]["z"].iloc[0] = 1.0
-    result = df.loc[(0, 0), "z"]
     if using_copy_on_write:
-        assert result == df_original.loc[0, "z"]
+        with pytest.raises(ChainedAssignmentError):
+            df.loc[0]["z"].iloc[0] = 1.0
+        assert df.loc[(0, 0), "z"] == df_original.loc[0, "z"]
     else:
+        df.loc[0]["z"].iloc[0] = 1.0
+        result = df.loc[(0, 0), "z"]
         assert result == 1
 
     # correct setting

pandas/tests/indexing/multiindex/test_partial.py

@@ -1,6 +1,7 @@
 import numpy as np
 import pytest
 
+from pandas.errors import ChainedAssignmentError
 import pandas.util._test_decorators as td
 
 from pandas import (
@@ -132,20 +133,26 @@ def test_partial_set(
         exp.iloc[65:85] = 0
         tm.assert_frame_equal(df, exp)
 
-        df["A"].loc[2000, 4] = 1
-        if not using_copy_on_write:
-            exp["A"].loc[2000, 4].values[:] = 1
+        if using_copy_on_write:
+            with pytest.raises(ChainedAssignmentError):
+                df["A"].loc[2000, 4] = 1
+            df.loc[(2000, 4), "A"] = 1
+        else:
+            df["A"].loc[2000, 4] = 1
+        exp.iloc[65:85, 0] = 1
         tm.assert_frame_equal(df, exp)
 
         df.loc[2000] = 5
         exp.iloc[:100] = 5
         tm.assert_frame_equal(df, exp)
 
         # this works...for now
-        df["A"].iloc[14] = 5
         if using_copy_on_write:
+            with pytest.raises(ChainedAssignmentError):
+                df["A"].iloc[14] = 5
             df["A"].iloc[14] == exp["A"].iloc[14]
         else:
+            df["A"].iloc[14] = 5
             assert df["A"].iloc[14] == 5
 
     @pytest.mark.parametrize("dtype", [int, float])

pandas/tests/indexing/multiindex/test_setitem.py

@@ -1,7 +1,10 @@
 import numpy as np
 import pytest
 
-from pandas.errors import SettingWithCopyError
+from pandas.errors import (
+    ChainedAssignmentError,
+    SettingWithCopyError,
+)
 import pandas.util._test_decorators as td
 
 import pandas as pd
@@ -501,8 +504,8 @@ def test_frame_setitem_copy_raises(
     # will raise/warn as its chained assignment
     df = multiindex_dataframe_random_data.T
     if using_copy_on_write:
-        # TODO(CoW) it would be nice if this could still warn/raise
-        df["foo"]["one"] = 2
+        with pytest.raises(ChainedAssignmentError):
+            df["foo"]["one"] = 2
     else:
         msg = "A value is trying to be set on a copy of a slice from a DataFrame"
         with pytest.raises(SettingWithCopyError, match=msg):
@@ -516,7 +519,8 @@ def test_frame_setitem_copy_no_write(
     expected = frame
     df = frame.copy()
     if using_copy_on_write:
-        df["foo"]["one"] = 2
+        with pytest.raises(ChainedAssignmentError):
+            df["foo"]["one"] = 2
     else:
         msg = "A value is trying to be set on a copy of a slice from a DataFrame"
         with pytest.raises(SettingWithCopyError, match=msg):