Show invalid comparator in error message

pola-rs · Aug 3, 2023 · 36ae804 · 36ae804
1 parent 66ee4f9
commit 36ae804
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 19 deletions.
diff --git a/py-polars/polars/lazyframe/frame.py b/py-polars/polars/lazyframe/frame.py
@@ -660,15 +660,26 @@ def __bool__(self) -> NoReturn:
             "cannot be used in boolean context with and/or/not operators. "
         )
 
-    def _comparison_error(self, other: Any) -> NoReturn:
-        raise TypeError("Cannot compare LazyFrames.")
-
-    __eq__ = _comparison_error
-    __ne__ = _comparison_error
-    __gt__ = _comparison_error
-    __lt__ = _comparison_error
-    __ge__ = _comparison_error
-    __le__ = _comparison_error
+    def _comparison_error(self, operator: str) -> NoReturn:
+        raise TypeError(f'"{operator}" comparison not supported for LazyFrame objects.')
+
+    def __eq__(self, other: Any) -> NoReturn:
+        self._comparison_error("==")
+
+    def __ne__(self, other: Any) -> NoReturn:
+        self._comparison_error("!=")
+
+    def __gt__(self, other: Any) -> NoReturn:
+        self._comparison_error(">")
+
+    def __lt__(self, other: Any) -> NoReturn:
+        self._comparison_error("<")
+
+    def __ge__(self, other: Any) -> NoReturn:
+        self._comparison_error(">=")
+
+    def __le__(self, other: Any) -> NoReturn:
+        self._comparison_error("<=")
 
     def __contains__(self, key: str) -> bool:
         return key in self.columns

diff --git a/py-polars/tests/unit/test_lazy.py b/py-polars/tests/unit/test_lazy.py
@@ -1462,22 +1462,22 @@ def test_compare_aggregation_between_lazy_and_eager_6904(
 
 
 @pytest.mark.parametrize(
-    "compare",
+    "comparators",
     [
-        pl.LazyFrame.__eq__,
-        pl.LazyFrame.__ne__,
-        pl.LazyFrame.__gt__,
-        pl.LazyFrame.__lt__,
-        pl.LazyFrame.__ge__,
-        pl.LazyFrame.__le__,
+        ("==", pl.LazyFrame.__eq__),
+        ("!=", pl.LazyFrame.__ne__),
+        (">", pl.LazyFrame.__gt__),
+        ("<", pl.LazyFrame.__lt__),
+        (">=", pl.LazyFrame.__ge__),
+        ("<=", pl.LazyFrame.__le__),
     ],
 )
 def test_lazy_comparison_operators(
-    compare: Callable[[pl.LazyFrame, pl.LazyFrame], NoReturn]
+    comparators: Callable[[pl.LazyFrame, pl.LazyFrame], NoReturn]
 ) -> None:
     # we cannot compare lazy frames, so all should raise a TypeError
     with pytest.raises(
         TypeError,
-        match="Cannot compare LazyFrames.",
+        match=f'"{comparators[0]}" comparison not supported for LazyFrame objects',
     ):
-        compare(pl.LazyFrame(), pl.LazyFrame())
+        comparators[1](pl.LazyFrame(), pl.LazyFrame())