Add comparison methods

jvadair · May 8, 2023 · 3ff978b · 3ff978b
1 parent 447ea9c
commit 3ff978b
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 0 deletions.
diff --git a/pyntree/__init__.py b/pyntree/__init__.py
@@ -233,3 +233,23 @@ def __imod__(self, other):
     def __ipow__(self, other):
         self.file.data[self.path[0]] **= other
         return self()
+
+    # Comparison methods (<, >, <=, >=, ==, !=)
+    def __lt__(self, other):
+        return True if self() < other() else False
+
+    def __le__(self, other):
+        return True if self() <= other() else False
+
+    def __gt__(self, other):
+        return True if self() > other() else False
+
+    def __ge__(self, other):
+        return True if self() >= other() else False
+
+    def __eq__(self, other):
+        return True if self() == other() else False
+
+    def __ne__(self, other):
+        return True if self() != other() else False
+
diff --git a/tests/tests.py b/tests/tests.py
@@ -351,6 +351,17 @@ def test_imul_str(self):
         db.a *= 3
         self.assertEqual(db.a(), 'aaa')
 
+    def test_cmp(self):
+        db = Node()
+        db.a = 1
+        db.b = 2
+        self.assertTrue(db.a < db.b)
+        self.assertTrue(db.a <= db.b)
+        self.assertFalse(db.a > db.b)
+        self.assertFalse(db.a >= db.b)
+        self.assertFalse(db.a == db.b)
+        self.assertTrue(db.a != db.b)
+
 
 
 if __name__ == '__main__':