Skip to content

Commit 640dad1

Browse files
majiayu000claude
andcommitted
TST: Add tests for logical operators in pd.col Expression class
Add comprehensive tests for the new logical operators (__and__, __or__, __xor__, __invert__) in the Expression class: - test_col_logical_ops: Parametrized tests for &, |, ^, ~ operators with both column-to-column and column-to-scalar operations - test_col_logical_ops_combined: Integration test demonstrating real-world usage in DataFrame.loc[] with combined conditions 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
1 parent 419fab8 commit 640dad1

File tree

1 file changed

+90
-0
lines changed

1 file changed

+90
-0
lines changed

pandas/tests/test_col.py

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,3 +97,93 @@ def mean(self):
9797
result = df.assign(b=pd.col("a").xyz.mean())
9898
expected = pd.DataFrame({"a": [1, 2, 3], "b": [2.0, 2.0, 2.0]})
9999
tm.assert_frame_equal(result, expected)
100+
101+
102+
@pytest.mark.parametrize(
103+
("expr", "expected_values", "expected_str"),
104+
[
105+
# __and__ and __rand__
106+
# a = [True, False, True, False], b = [False, True, True, True]
107+
# a & b = [False, False, True, False]
108+
(
109+
pd.col("a") & pd.col("b"),
110+
[False, False, True, False],
111+
"(col('a') & col('b'))",
112+
),
113+
(
114+
pd.col("a") & True,
115+
[True, False, True, False],
116+
"(col('a') & True)",
117+
),
118+
# __or__ and __ror__
119+
(
120+
pd.col("a") | pd.col("b"),
121+
[True, True, True, True],
122+
"(col('a') | col('b'))",
123+
),
124+
(
125+
pd.col("a") | False,
126+
[True, False, True, False],
127+
"(col('a') | False)",
128+
),
129+
# __xor__ and __rxor__
130+
# a ^ b = [True, True, False, True]
131+
(
132+
pd.col("a") ^ pd.col("b"),
133+
[True, True, False, True],
134+
"(col('a') ^ col('b'))",
135+
),
136+
(
137+
pd.col("a") ^ True,
138+
[False, True, False, True],
139+
"(col('a') ^ True)",
140+
),
141+
# __invert__
142+
(
143+
~pd.col("a"),
144+
[False, True, False, True],
145+
"(~col('a'))",
146+
),
147+
],
148+
)
149+
def test_col_logical_ops(
150+
expr: Expression, expected_values: list[object], expected_str: str
151+
) -> None:
152+
df = pd.DataFrame({"a": [True, False, True, False], "b": [False, True, True, True]})
153+
result = df.assign(c=expr)
154+
expected = pd.DataFrame(
155+
{
156+
"a": [True, False, True, False],
157+
"b": [False, True, True, True],
158+
"c": expected_values,
159+
}
160+
)
161+
tm.assert_frame_equal(result, expected)
162+
assert str(expr) == expected_str
163+
164+
165+
def test_col_logical_ops_combined() -> None:
166+
"""Test combining multiple logical operators like in DataFrame.loc[]."""
167+
df = pd.DataFrame(
168+
{
169+
"x": [1, 2, 3, 4, 5],
170+
"y": [10, 20, 30, 40, 50],
171+
}
172+
)
173+
# Test combining conditions with & operator
174+
expr = (pd.col("x") > 2) & (pd.col("y") < 45)
175+
result = df.loc[expr]
176+
expected = df.loc[(df["x"] > 2) & (df["y"] < 45)]
177+
tm.assert_frame_equal(result, expected)
178+
179+
# Test combining conditions with | operator
180+
expr = (pd.col("x") == 1) | (pd.col("x") == 5)
181+
result = df.loc[expr]
182+
expected = df.loc[(df["x"] == 1) | (df["x"] == 5)]
183+
tm.assert_frame_equal(result, expected)
184+
185+
# Test negation with ~
186+
expr = ~(pd.col("x") > 3)
187+
result = df.loc[expr]
188+
expected = df.loc[~(df["x"] > 3)]
189+
tm.assert_frame_equal(result, expected)

0 commit comments

Comments
 (0)