68 lines
2.0 KiB
Python
68 lines
2.0 KiB
Python
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)
|