| 1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071 |
- import numpy as np
- from pandas import (
- DataFrame,
- MultiIndex,
- Series,
- )
- import pandas._testing as tm
- class TestDataFramePop:
- def test_pop(self, float_frame):
- float_frame.columns.name = "baz"
- float_frame.pop("A")
- assert "A" not in float_frame
- float_frame["foo"] = "bar"
- float_frame.pop("foo")
- assert "foo" not in float_frame
- assert float_frame.columns.name == "baz"
- # gh-10912: inplace ops cause caching issue
- a = DataFrame([[1, 2, 3], [4, 5, 6]], columns=["A", "B", "C"], index=["X", "Y"])
- b = a.pop("B")
- b += 1
- # original frame
- expected = DataFrame([[1, 3], [4, 6]], columns=["A", "C"], index=["X", "Y"])
- tm.assert_frame_equal(a, expected)
- # result
- expected = Series([2, 5], index=["X", "Y"], name="B") + 1
- tm.assert_series_equal(b, expected)
- def test_pop_non_unique_cols(self):
- df = DataFrame({0: [0, 1], 1: [0, 1], 2: [4, 5]})
- df.columns = ["a", "b", "a"]
- res = df.pop("a")
- assert type(res) == DataFrame
- assert len(res) == 2
- assert len(df.columns) == 1
- assert "b" in df.columns
- assert "a" not in df.columns
- assert len(df.index) == 2
- def test_mixed_depth_pop(self):
- arrays = [
- ["a", "top", "top", "routine1", "routine1", "routine2"],
- ["", "OD", "OD", "result1", "result2", "result1"],
- ["", "wx", "wy", "", "", ""],
- ]
- tuples = sorted(zip(*arrays))
- index = MultiIndex.from_tuples(tuples)
- df = DataFrame(np.random.randn(4, 6), columns=index)
- df1 = df.copy()
- df2 = df.copy()
- result = df1.pop("a")
- expected = df2.pop(("a", "", ""))
- tm.assert_series_equal(expected, result, check_names=False)
- tm.assert_frame_equal(df1, df2)
- assert result.name == "a"
- expected = df1["top"]
- df1 = df1.drop(["top"], axis=1)
- result = df2.pop("top")
- tm.assert_frame_equal(expected, result)
- tm.assert_frame_equal(df1, df2)
|
