channel: round out the rich comparison operators

lumicks · Oct 31, 2024 · 2a24ceb · 2a24ceb
1 parent c672518
commit 2a24ceb
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 1 deletion.
diff --git a/lumicks/pylake/channel.py b/lumicks/pylake/channel.py
@@ -170,6 +170,42 @@ def __rmul__(self, other):
             labels=self._generate_labels(other, "*", self, keep_unit=False),
         )
 
+    def __eq__(self, other):
+        return Slice(
+            self._src._with_data(self.data == self._unpack_other(other)),
+            labels=self._generate_labels(self, "==", other, keep_unit=False),
+        )
+
+    def __lt__(self, other):
+        return Slice(
+            self._src._with_data(self.data < self._unpack_other(other)),
+            labels=self._generate_labels(self, "<", other, keep_unit=False),
+        )
+
+    def __le__(self, other):
+        return Slice(
+            self._src._with_data(self.data <= self._unpack_other(other)),
+            labels=self._generate_labels(self, "<=", other, keep_unit=False),
+        )
+
+    def __gt__(self, other):
+        return Slice(
+            self._src._with_data(self.data > self._unpack_other(other)),
+            labels=self._generate_labels(self, ">", other, keep_unit=False),
+        )
+
+    def __ge__(self, other):
+        return Slice(
+            self._src._with_data(self.data >= self._unpack_other(other)),
+            labels=self._generate_labels(self, ">=", other, keep_unit=False),
+        )
+
+    def __ne__(self, other):
+        return Slice(
+            self._src._with_data(self.data != self._unpack_other(other)),
+            labels=self._generate_labels(self, "!=", other, keep_unit=False),
+        )
+
     def __array__(self, *args, **kwargs):
         return np.asarray(self.data, *args, **kwargs)
 

diff --git a/lumicks/pylake/tests/test_channels/test_arithmetic.py b/lumicks/pylake/tests/test_channels/test_arithmetic.py
@@ -27,6 +27,12 @@
     ["__rtruediv__", False, "/", False],
     ["__rmul__", False, "*", False],
     ["__rpow__", False, "**", False],
+    ["__eq__", False, "==", False],
+    ["__ne__", False, "!=", False],
+    ["__lt__", False, "<", False],
+    ["__le__", False, "<=", False],
+    ["__gt__", False, ">", False],
+    ["__ge__", False, ">=", False],
 ]
 
 
@@ -181,7 +187,7 @@ def test_negation(channel_slice):
 
 def test_negation_timetags_not_implemented():
     with pytest.raises(NotImplementedError):
-        negated_timetags = -timetags
+        _ = -timetags
 
 
 def test_labels_slices():