feat[next]: Add missing UnitRange comparison functions

src/gt4py/next/common.py

@@ -141,15 +141,68 @@ def __and__(self, other: Set[int]) -> UnitRange:
         else:
             raise NotImplementedError("Can only find the intersection between UnitRange instances.")
 
-    def __le__(self, other: Set[int]):
+    def __contains__(self, value: Any) -> bool:
+        if not isinstance(value, core_defs.INTEGRAL_TYPES):
+            return False
+        if value == Infinity.positive() or value == Infinity.negative():
+            # raising error was an adhoc decision, feel free to improve
+            raise ValueError("Cannot check if Infinity is in UnitRange.")
+        return value >= self.start and value < self.stop
+
+    def __le__(self, other: Set[int]) -> bool:
         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 __lt__(self, other: Set[int]) -> bool:
+        if isinstance(other, UnitRange):
+            return (self.start > other.start and self.stop <= other.stop) or (
+                self.start >= other.start and self.stop < other.stop
+            )
+        elif len(self) == Infinity.positive():
+            return False
+        else:
+            return Set.__lt__(self, other)
+
+    def __ge__(self, other: Set[int]) -> bool:
+        if isinstance(other, UnitRange):
+            return self.start <= other.start and self.stop >= other.stop
+        elif len(self) == Infinity.positive():
+            for v in other:
+                if v not in self:
+                    return False
+            return True
+        else:
+            return Set.__ge__(self, other)
+
+    def __gt__(self, other: Set[int]) -> bool:
+        if isinstance(other, UnitRange):
+            return (self.start < other.start and self.stop >= other.stop) or (
+                self.start <= other.start and self.stop > other.stop
+            )
+        elif len(self) == Infinity.positive():
+            for v in other:
+                if v not in self:
+                    return False
+            return True
+        else:
+            return Set.__gt__(self, other)
+
+    def __eq__(self, other: Any) -> bool:
+        if isinstance(other, UnitRange):
+            return self.start == other.start and self.stop == other.stop
+        elif len(self) == Infinity.positive():
+            return False
+        elif isinstance(other, Set):
+            return Set.__eq__(self, other)
+        else:
+            return False
+
+    def __ne__(self, other: Any) -> bool:
+        return not self.__eq__(other)
 
     def __str__(self) -> str:
         return f"({self.start}:{self.stop})"

tests/next_tests/unit_tests/test_common.py

@@ -35,6 +35,11 @@
 KDim = Dimension("KDim", kind=DimensionKind.VERTICAL)
 
 
+@pytest.fixture(params=[Infinity.positive(), Infinity.negative()])
+def inf(request):
+    yield request.param
+
+
 @pytest.fixture
 def a_domain():
     return Domain((IDim, UnitRange(0, 10)), (JDim, UnitRange(5, 15)), (KDim, UnitRange(20, 30)))
@@ -151,6 +156,61 @@ def test_mixed_infinity_range():
     assert len(mixed_inf_range) == Infinity.positive()
 
 
+def test_range_contains():
+    assert 1 in UnitRange(0, 2)
+    assert 1 not in UnitRange(0, 1)
+    assert 1 in UnitRange(0, Infinity.positive())
+    assert 1 in UnitRange(Infinity.negative(), 2)
+    assert 1 in UnitRange(Infinity.negative(), Infinity.positive())
+    assert "s" not in UnitRange(Infinity.negative(), Infinity.positive())
+
+
+def test_range_contains_infinity(inf):
+    with pytest.raises(ValueError):
+        inf in UnitRange(Infinity.negative(), 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),
+        (operator.ge, UnitRange(-2, 3), UnitRange(-1, 2), True),
+        (operator.ge, UnitRange(-2, 3), {-2, -1, 0, 1, 2}, True),
+        (operator.ge, UnitRange(-2, 3), {-2, -1, 0, 1, 2, 3}, False),
+        (operator.ge, UnitRange(Infinity.negative(), 3), UnitRange(Infinity.negative(), 2), True),
+        (operator.ge, UnitRange(Infinity.negative(), 3), {1, 2}, True),
+        (operator.lt, UnitRange(-1, 2), UnitRange(-2, 2), True),
+        (operator.lt, UnitRange(-2, 1), UnitRange(-2, 2), True),
+        (operator.lt, UnitRange(-2, 2), {-1, 0, 1, 2}, False),
+        (operator.lt, UnitRange(-2, 2), {-2, -1, 0, 1, 2}, True),
+        (operator.lt, UnitRange(-2, 2), {-3, -2, -1, 0, 1, 2}, True),
+        (operator.lt, UnitRange(Infinity.negative(), 2), UnitRange(Infinity.negative(), 3), True),
+        (operator.lt, UnitRange(Infinity.negative(), 2), {1, 2, 3}, False),
+        (operator.gt, UnitRange(-2, 2), UnitRange(-1, 2), True),
+        (operator.gt, UnitRange(-2, 2), UnitRange(-2, 1), True),
+        (operator.gt, UnitRange(-2, 2), {-1, 0, 1}, True),
+        (operator.gt, UnitRange(-2, 2), {-2, -1, 0}, True),
+        (operator.gt, UnitRange(-2, 2), {-2, -1, 0, 1}, False),
+        (operator.gt, UnitRange(Infinity.negative(), 3), UnitRange(Infinity.negative(), 2), True),
+        (operator.gt, UnitRange(Infinity.negative(), 2), {0, 1}, True),
+        (operator.gt, UnitRange(Infinity.negative(), 2), {1, 2}, False),
+        (operator.eq, UnitRange(Infinity.negative(), 2), UnitRange(Infinity.negative(), 2), True),
+        (operator.eq, UnitRange(-2, 2), {-2, -1, 0, 1}, True),
+        (operator.eq, UnitRange(-2, 2), {-2, 1}, False),
+        (operator.ne, UnitRange(Infinity.negative(), 2), UnitRange(Infinity.negative(), 3), True),
+        (operator.ne, UnitRange(Infinity.negative(), 2), UnitRange(Infinity.negative(), 2), False),
+        (operator.ne, UnitRange(-2, 2), {-2, -1, 0}, True),
+    ],
+)
+def test_range_comparison(op, rng1, rng2, expected):
+    assert op(rng1, rng2) == expected
+
+
 @pytest.mark.parametrize(
     "op, rng1, rng2, expected",
     [