1import pytest
2
3from pandas import DataFrame
4import pandas._testing as tm
5
6
7class TestAssign:
8    def test_assign(self):
9        df = DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]})
10        original = df.copy()
11        result = df.assign(C=df.B / df.A)
12        expected = df.copy()
13        expected["C"] = [4, 2.5, 2]
14        tm.assert_frame_equal(result, expected)
15
16        # lambda syntax
17        result = df.assign(C=lambda x: x.B / x.A)
18        tm.assert_frame_equal(result, expected)
19
20        # original is unmodified
21        tm.assert_frame_equal(df, original)
22
23        # Non-Series array-like
24        result = df.assign(C=[4, 2.5, 2])
25        tm.assert_frame_equal(result, expected)
26        # original is unmodified
27        tm.assert_frame_equal(df, original)
28
29        result = df.assign(B=df.B / df.A)
30        expected = expected.drop("B", axis=1).rename(columns={"C": "B"})
31        tm.assert_frame_equal(result, expected)
32
33        # overwrite
34        result = df.assign(A=df.A + df.B)
35        expected = df.copy()
36        expected["A"] = [5, 7, 9]
37        tm.assert_frame_equal(result, expected)
38
39        # lambda
40        result = df.assign(A=lambda x: x.A + x.B)
41        tm.assert_frame_equal(result, expected)
42
43    def test_assign_multiple(self):
44        df = DataFrame([[1, 4], [2, 5], [3, 6]], columns=["A", "B"])
45        result = df.assign(C=[7, 8, 9], D=df.A, E=lambda x: x.B)
46        expected = DataFrame(
47            [[1, 4, 7, 1, 4], [2, 5, 8, 2, 5], [3, 6, 9, 3, 6]], columns=list("ABCDE")
48        )
49        tm.assert_frame_equal(result, expected)
50
51    def test_assign_order(self):
52        # GH 9818
53        df = DataFrame([[1, 2], [3, 4]], columns=["A", "B"])
54        result = df.assign(D=df.A + df.B, C=df.A - df.B)
55
56        expected = DataFrame([[1, 2, 3, -1], [3, 4, 7, -1]], columns=list("ABDC"))
57        tm.assert_frame_equal(result, expected)
58        result = df.assign(C=df.A - df.B, D=df.A + df.B)
59
60        expected = DataFrame([[1, 2, -1, 3], [3, 4, -1, 7]], columns=list("ABCD"))
61
62        tm.assert_frame_equal(result, expected)
63
64    def test_assign_bad(self):
65        df = DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]})
66
67        # non-keyword argument
68        msg = r"assign\(\) takes 1 positional argument but 2 were given"
69        with pytest.raises(TypeError, match=msg):
70            df.assign(lambda x: x.A)
71        msg = "'DataFrame' object has no attribute 'C'"
72        with pytest.raises(AttributeError, match=msg):
73            df.assign(C=df.A, D=df.A + df.C)
74
75    def test_assign_dependent(self):
76        df = DataFrame({"A": [1, 2], "B": [3, 4]})
77
78        result = df.assign(C=df.A, D=lambda x: x["A"] + x["C"])
79        expected = DataFrame([[1, 3, 1, 2], [2, 4, 2, 4]], columns=list("ABCD"))
80        tm.assert_frame_equal(result, expected)
81
82        result = df.assign(C=lambda df: df.A, D=lambda df: df["A"] + df["C"])
83        expected = DataFrame([[1, 3, 1, 2], [2, 4, 2, 4]], columns=list("ABCD"))
84        tm.assert_frame_equal(result, expected)
85