Implement chained assignment detection with is_local_in_caller_frame approach for py3.14 compat (reverting cython SetitemMixin class) (#63186)

Co-authored-by: Matt Page <[email protected]>

Author: jorisvandenbossche

pandas/_libs/internals.pyi

@@ -94,7 +94,3 @@ class BlockValuesRefs:
     def add_reference(self, blk: Block) -> None: ...
     def add_index_reference(self, index: Index) -> None: ...
     def has_reference(self) -> bool: ...
-
-class SetitemMixin:
-    def __setitem__(self, key, value) -> None: ...
-    def __delitem__(self, key) -> None: ...

pandas/_libs/internals.pyx

@@ -1,9 +1,6 @@
 from collections import defaultdict
-import sys
-import warnings
 
 cimport cython
-from cpython cimport PY_VERSION_HEX
 from cpython.object cimport PyObject
 from cpython.pyport cimport PY_SSIZE_T_MAX
 from cpython.slice cimport PySlice_GetIndicesEx
@@ -23,9 +20,6 @@ from numpy cimport (
 cnp.import_array()
 
 from pandas._libs.algos import ensure_int64
-from pandas.compat import CHAINED_WARNING_DISABLED
-from pandas.errors import ChainedAssignmentError
-from pandas.errors.cow import _chained_assignment_msg
 
 from pandas._libs.util cimport (
     is_array,
@@ -1002,47 +996,3 @@ cdef class BlockValuesRefs:
                 return self._has_reference_maybe_locked()
         ELSE:
             return self._has_reference_maybe_locked()
-
-
-cdef extern from "Python.h":
-    """
-    // python version < 3.14
-    #if PY_VERSION_HEX < 0x030E0000
-    // This function is unused and is declared to avoid a build warning
-    int __Pyx_PyUnstable_Object_IsUniqueReferencedTemporary(PyObject *ref) {
-        return Py_REFCNT(ref) == 1;
-    }
-    #else
-    #define __Pyx_PyUnstable_Object_IsUniqueReferencedTemporary \
-        PyUnstable_Object_IsUniqueReferencedTemporary
-    #endif
-    """
-    int PyUnstable_Object_IsUniqueReferencedTemporary\
-        "__Pyx_PyUnstable_Object_IsUniqueReferencedTemporary"(object o) except -1
-
-
-# Python version compatibility for PyUnstable_Object_IsUniqueReferencedTemporary
-cdef inline bint _is_unique_referenced_temporary(object obj) except -1:
-    if PY_VERSION_HEX >= 0x030E0000:
-        # Python 3.14+ has PyUnstable_Object_IsUniqueReferencedTemporary
-        return PyUnstable_Object_IsUniqueReferencedTemporary(obj)
-    else:
-        # Fallback for older Python versions using sys.getrefcount
-        return sys.getrefcount(obj) <= 1
-
-
-cdef class SetitemMixin:
-    # class used in DataFrame and Series for checking for chained assignment
-
-    def __setitem__(self, key, value) -> None:
-        cdef bint is_unique = 0
-        if not CHAINED_WARNING_DISABLED:
-            is_unique = _is_unique_referenced_temporary(self)
-            if is_unique:
-                warnings.warn(
-                    _chained_assignment_msg, ChainedAssignmentError, stacklevel=1
-                )
-        self._setitem(key, value)
-
-    def __delitem__(self, key) -> None:
-        self._delitem(key)

pandas/_testing/contexts.py

@@ -12,10 +12,7 @@
 )
 import uuid
 
-from pandas.compat import (
-    CHAINED_WARNING_DISABLED,
-    CHAINED_WARNING_DISABLED_INPLACE_METHOD,
-)
+from pandas.compat import CHAINED_WARNING_DISABLED
 from pandas.errors import ChainedAssignmentError
 
 from pandas.io.common import get_handle
@@ -163,18 +160,10 @@ def with_csv_dialect(name: str, **kwargs) -> Generator[None]:
         csv.unregister_dialect(name)
 
 
-def raises_chained_assignment_error(
-    extra_warnings=(), extra_match=(), inplace_method=False
-):
+def raises_chained_assignment_error(extra_warnings=(), extra_match=()):
     from pandas._testing import assert_produces_warning
 
-    WARNING_DISABLED = (
-        CHAINED_WARNING_DISABLED_INPLACE_METHOD
-        if inplace_method
-        else CHAINED_WARNING_DISABLED
-    )
-
-    if WARNING_DISABLED:
+    if CHAINED_WARNING_DISABLED:
         if not extra_warnings:
             from contextlib import nullcontext
 

pandas/compat/__init__.py

@@ -17,7 +17,6 @@
 
 from pandas.compat._constants import (
     CHAINED_WARNING_DISABLED,
-    CHAINED_WARNING_DISABLED_INPLACE_METHOD,
     IS64,
     ISMUSL,
     PY312,
@@ -154,7 +153,6 @@ def is_ci_environment() -> bool:
 
 __all__ = [
     "CHAINED_WARNING_DISABLED",
-    "CHAINED_WARNING_DISABLED_INPLACE_METHOD",
     "HAS_PYARROW",
     "IS64",
     "ISMUSL",

pandas/compat/_constants.py

@@ -18,9 +18,11 @@
 PYPY = platform.python_implementation() == "PyPy"
 WASM = (sys.platform == "emscripten") or (platform.machine() in ["wasm32", "wasm64"])
 ISMUSL = "musl" in (sysconfig.get_config_var("HOST_GNU_TYPE") or "")
-REF_COUNT = 2
-CHAINED_WARNING_DISABLED = PYPY or (PY314 and not sys._is_gil_enabled())  # type: ignore[attr-defined]
-CHAINED_WARNING_DISABLED_INPLACE_METHOD = PYPY or PY314
+# the refcount for self in a chained __setitem__/.(i)loc indexing/method call
+REF_COUNT = 2 if PY314 else 3
+REF_COUNT_IDX = 2
+REF_COUNT_METHOD = 1 if PY314 else 2
+CHAINED_WARNING_DISABLED = PYPY
 
 
 __all__ = [

pandas/compat/pickle_compat.py

@@ -22,7 +22,6 @@
     PeriodArray,
     TimedeltaArray,
 )
-from pandas.core.generic import NDFrame
 from pandas.core.internals import BlockManager
 
 if TYPE_CHECKING:
@@ -91,10 +90,6 @@ def load_reduce(self) -> None:
                 cls = args[0]
                 stack[-1] = NDArrayBacked.__new__(*args)
                 return
-            elif args and issubclass(args[0], NDFrame):
-                cls = args[0]
-                stack[-1] = cls.__new__(cls)
-                return
             raise
 
     dispatch[pickle.REDUCE[0]] = load_reduce  # type: ignore[assignment]

pandas/core/common.py

@@ -22,6 +22,7 @@
 import contextlib
 from functools import partial
 import inspect
+import sys
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -650,3 +651,29 @@ def fill_missing_names(names: Sequence[Hashable | None]) -> list[Hashable]:
         list of column names with the None values replaced.
     """
     return [f"level_{i}" if name is None else name for i, name in enumerate(names)]
+
+
+def is_local_in_caller_frame(obj):
+    """
+    Helper function used in detecting chained assignment.
+
+    If the pandas object (DataFrame/Series) is a local variable
+    in the caller's frame, it should not be a case of chained
+    assignment or method call.
+
+    For example:
+
+    def test():
+        df = pd.DataFrame(...)
+        df["a"] = 1  # not chained assignment
+
+    Inside ``df.__setitem__``, we call this function to check whether `df`
+    (`self`) is a local variable in `test` frame (the frame calling setitem). If
+    so, we know it is not a case of chained assignment (even when the refcount
+    of `df` is below the threshold due to optimization of local variables).
+    """
+    frame = sys._getframe(2)
+    for v in frame.f_locals.values():
+        if v is obj:
+            return True
+    return False

pandas/core/frame.py

@@ -48,11 +48,11 @@
     properties,
 )
 from pandas._libs.hashtable import duplicated
-from pandas._libs.internals import SetitemMixin
 from pandas._libs.lib import is_range_indexer
+from pandas.compat import CHAINED_WARNING_DISABLED
 from pandas.compat._constants import (
-    CHAINED_WARNING_DISABLED_INPLACE_METHOD,
     REF_COUNT,
+    REF_COUNT_METHOD,
 )
 from pandas.compat._optional import import_optional_dependency
 from pandas.compat.numpy import function as nv
@@ -63,6 +63,7 @@
 )
 from pandas.errors.cow import (
     _chained_assignment_method_msg,
+    _chained_assignment_msg,
 )
 from pandas.util._decorators import (
     Appender,
@@ -522,7 +523,7 @@
 
 
 @set_module("pandas")
-class DataFrame(SetitemMixin, NDFrame, OpsMixin):
+class DataFrame(NDFrame, OpsMixin):
     """
     Two-dimensional, size-mutable, potentially heterogeneous tabular data.
 
@@ -669,11 +670,6 @@ class DataFrame(SetitemMixin, NDFrame, OpsMixin):
     #  and ExtensionArray.  Should NOT be overridden by subclasses.
     __pandas_priority__ = 4000
 
-    # override those to avoid inheriting from SetitemMixin (cython generates
-    # them by default)
-    __reduce__ = object.__reduce__
-    __setstate__ = NDFrame.__setstate__
-
     @property
     def _constructor(self) -> type[DataFrame]:
         return DataFrame
@@ -4377,8 +4373,7 @@ def isetitem(self, loc, value) -> None:
         arraylike, refs = self._sanitize_column(value)
         self._iset_item_mgr(loc, arraylike, inplace=False, refs=refs)
 
-    # def __setitem__() is implemented in SetitemMixin and dispatches to this method
-    def _setitem(self, key, value) -> None:
+    def __setitem__(self, key, value) -> None:
         """
         Set item(s) in DataFrame by key.
 
@@ -4462,6 +4457,14 @@ def _setitem(self, key, value) -> None:
         z  3  50
         # Values for 'a' and 'b' are completely ignored!
         """
+        if not CHAINED_WARNING_DISABLED:
+            if sys.getrefcount(self) <= REF_COUNT and not com.is_local_in_caller_frame(
+                self
+            ):
+                warnings.warn(
+                    _chained_assignment_msg, ChainedAssignmentError, stacklevel=2
+                )
+
         key = com.apply_if_callable(key, self)
 
         # see if we can slice the rows
@@ -9414,8 +9417,10 @@ def update(
         1  2  500.0
         2  3    6.0
         """
-        if not CHAINED_WARNING_DISABLED_INPLACE_METHOD:
-            if sys.getrefcount(self) <= REF_COUNT:
+        if not CHAINED_WARNING_DISABLED:
+            if sys.getrefcount(
+                self
+            ) <= REF_COUNT_METHOD and not com.is_local_in_caller_frame(self):
                 warnings.warn(
                     _chained_assignment_method_msg,
                     ChainedAssignmentError,