diff --git a/pandas/tests/copy_view/test_methods.py b/pandas/tests/copy_view/test_methods.py index 6b54345723118..791a744e5b4e3 100644 --- a/pandas/tests/copy_view/test_methods.py +++ b/pandas/tests/copy_view/test_methods.py @@ -1,6 +1,9 @@ import numpy as np import pytest +from pandas.errors import SettingWithCopyWarning + +import pandas as pd from pandas import ( DataFrame, Index, @@ -1308,3 +1311,90 @@ def test_isetitem(using_copy_on_write): assert np.shares_memory(get_array(df, "c"), get_array(df2, "c")) else: assert not np.shares_memory(get_array(df, "c"), get_array(df2, "c")) + + +@pytest.mark.parametrize("key", ["a", ["a"]]) +def test_get(using_copy_on_write, key): + df = DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) + df_orig = df.copy() + + result = df.get(key) + + if using_copy_on_write: + assert np.shares_memory(get_array(result, "a"), get_array(df, "a")) + result.iloc[0] = 0 + assert not np.shares_memory(get_array(result, "a"), get_array(df, "a")) + tm.assert_frame_equal(df, df_orig) + else: + # for non-CoW it depends on whether we got a Series or DataFrame if it + # is a view or copy or triggers a warning or not + warn = SettingWithCopyWarning if isinstance(key, list) else None + with pd.option_context("chained_assignment", "warn"): + with tm.assert_produces_warning(warn): + result.iloc[0] = 0 + + if isinstance(key, list): + tm.assert_frame_equal(df, df_orig) + else: + assert df.iloc[0, 0] == 0 + + +@pytest.mark.parametrize("axis, key", [(0, 0), (1, "a")]) +@pytest.mark.parametrize( + "dtype", ["int64", "float64"], ids=["single-block", "mixed-block"] +) +def test_xs(using_copy_on_write, using_array_manager, axis, key, dtype): + single_block = (dtype == "int64") and not using_array_manager + is_view = single_block or (using_array_manager and axis == 1) + df = DataFrame( + {"a": [1, 2, 3], "b": [4, 5, 6], "c": np.array([7, 8, 9], dtype=dtype)} + ) + df_orig = df.copy() + + result = df.xs(key, axis=axis) + + if axis == 1 or single_block: + assert np.shares_memory(get_array(df, "a"), get_array(result)) + elif using_copy_on_write: + assert result._mgr._has_no_reference(0) + + if using_copy_on_write or is_view: + result.iloc[0] = 0 + else: + with pd.option_context("chained_assignment", "warn"): + with tm.assert_produces_warning(SettingWithCopyWarning): + result.iloc[0] = 0 + + if using_copy_on_write or (not single_block and axis == 0): + tm.assert_frame_equal(df, df_orig) + else: + assert df.iloc[0, 0] == 0 + + +@pytest.mark.parametrize("axis", [0, 1]) +@pytest.mark.parametrize("key, level", [("l1", 0), (2, 1)]) +def test_xs_multiindex(using_copy_on_write, using_array_manager, key, level, axis): + arr = np.arange(18).reshape(6, 3) + index = MultiIndex.from_product([["l1", "l2"], [1, 2, 3]], names=["lev1", "lev2"]) + df = DataFrame(arr, index=index, columns=list("abc")) + if axis == 1: + df = df.transpose().copy() + df_orig = df.copy() + + result = df.xs(key, level=level, axis=axis) + + if level == 0: + assert np.shares_memory( + get_array(df, df.columns[0]), get_array(result, result.columns[0]) + ) + + warn = ( + SettingWithCopyWarning + if not using_copy_on_write and not using_array_manager + else None + ) + with pd.option_context("chained_assignment", "warn"): + with tm.assert_produces_warning(warn): + result.iloc[0, 0] = 0 + + tm.assert_frame_equal(df, df_orig) diff --git a/pandas/tests/copy_view/util.py b/pandas/tests/copy_view/util.py index 8e53c1644b0a2..ba9b0ecdebb58 100644 --- a/pandas/tests/copy_view/util.py +++ b/pandas/tests/copy_view/util.py @@ -10,7 +10,7 @@ def get_array(obj, col=None): which triggers tracking references / CoW (and we might be testing that this is done by some other operation). """ - if isinstance(obj, Series) and (obj is None or obj.name == col): + if isinstance(obj, Series) and (col is None or obj.name == col): return obj._values assert col is not None icol = obj.columns.get_loc(col)