diff --git a/pandas/core/col.py b/pandas/core/col.py index 39c4a7fd016c2..b8e965845b4b6 100644 --- a/pandas/core/col.py +++ b/pandas/core/col.py @@ -37,6 +37,12 @@ "__lt__": "<", "__eq__": "==", "__ne__": "!=", + "__and__": "&", + "__rand__": "&", + "__or__": "|", + "__ror__": "|", + "__xor__": "^", + "__rxor__": "^", } @@ -157,6 +163,28 @@ def __mod__(self, other: Any) -> Expression: def __rmod__(self, other: Any) -> Expression: return self._with_binary_op("__rmod__", other) + # Logical ops + def __and__(self, other: Any) -> Expression: + return self._with_binary_op("__and__", other) + + def __rand__(self, other: Any) -> Expression: + return self._with_binary_op("__rand__", other) + + def __or__(self, other: Any) -> Expression: + return self._with_binary_op("__or__", other) + + def __ror__(self, other: Any) -> Expression: + return self._with_binary_op("__ror__", other) + + def __xor__(self, other: Any) -> Expression: + return self._with_binary_op("__xor__", other) + + def __rxor__(self, other: Any) -> Expression: + return self._with_binary_op("__rxor__", other) + + def __invert__(self) -> Expression: + return Expression(lambda df: ~self(df), f"(~{self._repr_str})") + def __array_ufunc__( self, ufunc: Callable[..., Any], method: str, *inputs: Any, **kwargs: Any ) -> Expression: diff --git a/pandas/tests/test_col.py b/pandas/tests/test_col.py index c884540abfed0..cf7901a912279 100644 --- a/pandas/tests/test_col.py +++ b/pandas/tests/test_col.py @@ -97,3 +97,65 @@ def mean(self): result = df.assign(b=pd.col("a").xyz.mean()) expected = pd.DataFrame({"a": [1, 2, 3], "b": [2.0, 2.0, 2.0]}) tm.assert_frame_equal(result, expected) + + +@pytest.mark.parametrize( + ("expr", "expected_values", "expected_str"), + [ + ( + pd.col("a") & pd.col("b"), + [False, False, True, False], + "(col('a') & col('b'))", + ), + ( + pd.col("a") & True, + [True, False, True, False], + "(col('a') & True)", + ), + ( + pd.col("a") | pd.col("b"), + [True, True, True, True], + "(col('a') | col('b'))", + ), + ( + pd.col("a") | False, + [True, False, True, False], + "(col('a') | False)", + ), + ( + pd.col("a") ^ pd.col("b"), + [True, True, False, True], + "(col('a') ^ col('b'))", + ), + ( + pd.col("a") ^ True, + [False, True, False, True], + "(col('a') ^ True)", + ), + ( + ~pd.col("a"), + [False, True, False, True], + "(~col('a'))", + ), + ], +) +def test_col_logical_ops( + expr: Expression, expected_values: list[bool], expected_str: str +) -> None: + # https://github.com/pandas-dev/pandas/issues/63322 + df = pd.DataFrame({"a": [True, False, True, False], "b": [False, True, True, True]}) + result = df.assign(c=expr) + expected = pd.DataFrame( + { + "a": [True, False, True, False], + "b": [False, True, True, True], + "c": expected_values, + } + ) + tm.assert_frame_equal(result, expected) + assert str(expr) == expected_str + + # Test that the expression works with .loc + result = df.loc[expr] + expected = df[expected_values] + tm.assert_frame_equal(result, expected)