Skip to content

Commit 7dee831

Browse files
Merge branch 'docstring' of https://github.com/zhangbowen-coder/pandas into docstring
2 parents cc41146 + 822f90b commit 7dee831

File tree

2 files changed

+82
-0
lines changed

2 files changed

+82
-0
lines changed

pandas/core/col.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,10 @@
3939
"__ne__": "!=",
4040
"__and__": "&",
4141
"__rand__": "&",
42+
"__or__": "|",
43+
"__ror__": "|",
44+
"__xor__": "^",
45+
"__rxor__": "^",
4246
}
4347

4448

@@ -159,12 +163,28 @@ def __mod__(self, other: Any) -> Expression:
159163
def __rmod__(self, other: Any) -> Expression:
160164
return self._with_binary_op("__rmod__", other)
161165

166+
# Logical ops
162167
def __and__(self, other: Any) -> Expression:
163168
return self._with_binary_op("__and__", other)
164169

165170
def __rand__(self, other: Any) -> Expression:
166171
return self._with_binary_op("__rand__", other)
167172

173+
def __or__(self, other: Any) -> Expression:
174+
return self._with_binary_op("__or__", other)
175+
176+
def __ror__(self, other: Any) -> Expression:
177+
return self._with_binary_op("__ror__", other)
178+
179+
def __xor__(self, other: Any) -> Expression:
180+
return self._with_binary_op("__xor__", other)
181+
182+
def __rxor__(self, other: Any) -> Expression:
183+
return self._with_binary_op("__rxor__", other)
184+
185+
def __invert__(self) -> Expression:
186+
return Expression(lambda df: ~self(df), f"(~{self._repr_str})")
187+
168188
def __array_ufunc__(
169189
self, ufunc: Callable[..., Any], method: str, *inputs: Any, **kwargs: Any
170190
) -> Expression:

pandas/tests/test_col.py

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -127,3 +127,65 @@ def mean(self):
127127
result = df.assign(b=pd.col("a").xyz.mean())
128128
expected = pd.DataFrame({"a": [1, 2, 3], "b": [2.0, 2.0, 2.0]})
129129
tm.assert_frame_equal(result, expected)
130+
131+
132+
@pytest.mark.parametrize(
133+
("expr", "expected_values", "expected_str"),
134+
[
135+
(
136+
pd.col("a") & pd.col("b"),
137+
[False, False, True, False],
138+
"(col('a') & col('b'))",
139+
),
140+
(
141+
pd.col("a") & True,
142+
[True, False, True, False],
143+
"(col('a') & True)",
144+
),
145+
(
146+
pd.col("a") | pd.col("b"),
147+
[True, True, True, True],
148+
"(col('a') | col('b'))",
149+
),
150+
(
151+
pd.col("a") | False,
152+
[True, False, True, False],
153+
"(col('a') | False)",
154+
),
155+
(
156+
pd.col("a") ^ pd.col("b"),
157+
[True, True, False, True],
158+
"(col('a') ^ col('b'))",
159+
),
160+
(
161+
pd.col("a") ^ True,
162+
[False, True, False, True],
163+
"(col('a') ^ True)",
164+
),
165+
(
166+
~pd.col("a"),
167+
[False, True, False, True],
168+
"(~col('a'))",
169+
),
170+
],
171+
)
172+
def test_col_logical_ops(
173+
expr: Expression, expected_values: list[bool], expected_str: str
174+
) -> None:
175+
# https://github.com/pandas-dev/pandas/issues/63322
176+
df = pd.DataFrame({"a": [True, False, True, False], "b": [False, True, True, True]})
177+
result = df.assign(c=expr)
178+
expected = pd.DataFrame(
179+
{
180+
"a": [True, False, True, False],
181+
"b": [False, True, True, True],
182+
"c": expected_values,
183+
}
184+
)
185+
tm.assert_frame_equal(result, expected)
186+
assert str(expr) == expected_str
187+
188+
# Test that the expression works with .loc
189+
result = df.loc[expr]
190+
expected = df[expected_values]
191+
tm.assert_frame_equal(result, expected)

0 commit comments

Comments
 (0)