implement le for UnitRange

src/gt4py/next/common.py

@@ -141,14 +141,16 @@ def __and__(self, other: Set[int]) -> UnitRange:
         else:
             raise NotImplementedError("Can only find the intersection between UnitRange instances.")
 
-    def __le__(self, other: Set[Any]):
+    def __le__(self, other: Set[int]):
         if isinstance(other, UnitRange):
             return self.start >= other.start and self.stop <= other.stop
         elif len(self) == Infinity.positive():
             return False
         else:
             return Set.__le__(self, other)
 
+    __ge__ = __lt__ = __gt__ = lambda self, other: NotImplemented
+
     def __str__(self) -> str:
         return f"({self.start}:{self.stop})"
 

tests/next_tests/unit_tests/test_common.py

@@ -11,6 +11,7 @@
 # distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
+import operator
 from typing import Optional, Pattern
 
 import pytest
@@ -150,6 +151,21 @@ def test_mixed_infinity_range():
     assert len(mixed_inf_range) == Infinity.positive()
 
 
+@pytest.mark.parametrize(
+    "op, rng1, rng2, expected",
+    [
+        (operator.le, UnitRange(-1, 2), UnitRange(-2, 3), True),
+        (operator.le, UnitRange(-1, 2), {-1, 0, 1}, True),
+        (operator.le, UnitRange(-1, 2), {-1, 0}, False),
+        (operator.le, UnitRange(-1, 2), {-2, -1, 0, 1, 2}, True),
+        (operator.le, UnitRange(Infinity.negative(), 2), UnitRange(Infinity.negative(), 3), True),
+        (operator.le, UnitRange(Infinity.negative(), 2), {1, 2, 3}, False),
+    ],
+)
+def test_range_comparison(op, rng1, rng2, expected):
+    assert op(rng1, rng2) == expected
+
+
 @pytest.mark.parametrize(
     "named_rng_like",
     [