Add comparison operators to Record class (#44)

Author: Schamper

dissect/esedb/record.py

@@ -40,6 +40,62 @@ def __init__(self, table: Table, node: Node):
         self._node = node
         self._data = RecordData(table, node)
 
+    def __str__(self) -> str:
+        column_values = serialise_record_column_values(self, max_columns=None)
+        return f"<Record {column_values}>"
+
+    def __repr__(self) -> str:
+        column_values = serialise_record_column_values(self)
+        return f"<Record {column_values}>"
+
+    def __hash__(self) -> int:
+        return hash((self._table, self._node))
+
+    def __getitem__(self, attr: str) -> RecordValue:
+        return self.get(attr)
+
+    def __getattr__(self, attr: str) -> RecordValue:
+        try:
+            return self.get(attr)
+        except KeyError:
+            return object.__getattribute__(self, attr)
+
+    def __eq__(self, value: object) -> bool:
+        if not isinstance(value, Record):
+            return False
+
+        return self._node.key == value._node.key
+
+    def __ne__(self, value: object) -> bool:
+        if not isinstance(value, Record):
+            return False
+
+        return self._node.key != value._node.key
+
+    def __lt__(self, value: object) -> bool:
+        if not isinstance(value, Record):
+            return False
+
+        return self._node.key < value._node.key
+
+    def __le__(self, value: object) -> bool:
+        if not isinstance(value, Record):
+            return False
+
+        return self._node.key <= value._node.key
+
+    def __gt__(self, value: object) -> bool:
+        if not isinstance(value, Record):
+            return False
+
+        return self._node.key > value._node.key
+
+    def __ge__(self, value: object) -> bool:
+        if not isinstance(value, Record):
+            return False
+
+        return self._node.key >= value._node.key
+
     def get(self, attr: str, raw: bool = False) -> RecordValue:
         """Retrieve a value from the record with the given name.
 
@@ -55,23 +111,6 @@ def get(self, attr: str, raw: bool = False) -> RecordValue:
     def as_dict(self, raw: bool = False) -> dict[str, RecordValue]:
         return self._data.as_dict(raw)
 
-    def __getitem__(self, attr: str) -> RecordValue:
-        return self.get(attr)
-
-    def __getattr__(self, attr: str) -> RecordValue:
-        try:
-            return self.get(attr)
-        except KeyError:
-            return object.__getattribute__(self, attr)
-
-    def __str__(self) -> str:
-        column_values = serialise_record_column_values(self, max_columns=None)
-        return f"<Record {column_values}>"
-
-    def __repr__(self) -> str:
-        column_values = serialise_record_column_values(self)
-        return f"<Record {column_values}>"
-
 
 class RecordData:
     """Record class for parsing and interacting with the on-disk record format.

tests/test_record.py

@@ -3,6 +3,7 @@
 from typing import BinaryIO
 
 from dissect.esedb.esedb import EseDB
+from dissect.esedb.record import Record
 
 
 def test_as_dict(basic_db: BinaryIO) -> None:
@@ -40,3 +41,25 @@ def test_as_dict(basic_db: BinaryIO) -> None:
             "DateTime": -4537072128574357504,
         },
     ]
+
+
+def test_comparison(basic_db: BinaryIO) -> None:
+    db = EseDB(basic_db)
+    table = db.table("basic")
+
+    records = list(table.records())
+    assert len(records) == 2
+
+    assert records[0] == records[0]
+    assert records[0] != records[1]
+
+    obj = Record(table, records[0]._node)
+    assert records[0] == obj
+    assert records[0] is not obj
+
+    assert records[0] < records[1]
+    assert records[0] <= records[1]
+    assert records[0] <= records[0]
+
+    assert set(records) == {records[0], records[1]}
+    assert set(records) | {obj} == {records[0], records[1]}