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..17f26bbfd4762 100644 --- a/pandas/tests/test_col.py +++ b/pandas/tests/test_col.py @@ -97,3 +97,93 @@ 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"), + [ + # __and__ and __rand__ + # a = [True, False, True, False], b = [False, True, True, True] + # a & b = [False, False, True, False] + ( + 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)", + ), + # __or__ and __ror__ + ( + 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)", + ), + # __xor__ and __rxor__ + # a ^ b = [True, True, False, True] + ( + 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)", + ), + # __invert__ + ( + ~pd.col("a"), + [False, True, False, True], + "(~col('a'))", + ), + ], +) +def test_col_logical_ops( + expr: Expression, expected_values: list[object], expected_str: str +) -> None: + 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 + + +def test_col_logical_ops_combined() -> None: + """Test combining multiple logical operators like in DataFrame.loc[].""" + df = pd.DataFrame( + { + "x": [1, 2, 3, 4, 5], + "y": [10, 20, 30, 40, 50], + } + ) + # Test combining conditions with & operator + expr = (pd.col("x") > 2) & (pd.col("y") < 45) + result = df.loc[expr] + expected = df.loc[(df["x"] > 2) & (df["y"] < 45)] + tm.assert_frame_equal(result, expected) + + # Test combining conditions with | operator + expr = (pd.col("x") == 1) | (pd.col("x") == 5) + result = df.loc[expr] + expected = df.loc[(df["x"] == 1) | (df["x"] == 5)] + tm.assert_frame_equal(result, expected) + + # Test negation with ~ + expr = ~(pd.col("x") > 3) + result = df.loc[expr] + expected = df.loc[~(df["x"] > 3)] + tm.assert_frame_equal(result, expected)