make Annotations sortable

src/pytorch_ie/core/document.py

@@ -191,6 +191,15 @@ def non_comparison_fields_and_values(self) -> Tuple[Tuple[str, Any], ...]:
             (f.name, getattr(self, f.name)) for f in dataclasses.fields(self) if not f.compare
         )
 
+    @property
+    def comparison_fields_and_values(self) -> Tuple[Tuple[str, Any], ...]:
+        # note that we exclude the special _targets field here
+        return tuple(
+            (f.name, getattr(self, f.name))
+            for f in dataclasses.fields(self)
+            if f.compare and f.name != "_targets"
+        )
+
     @property
     def _id(self) -> int:
         # also calculate the hash over the non-comparison fields for id creation because otherwise
@@ -361,6 +370,21 @@ def copy_with_store(
 
         return self.copy(**overrides)
 
+    def __lt__(self, other: "Annotation") -> bool:
+        # sort only by comparison fields
+        for (name, value), (other_name, other_value) in zip(
+            self.comparison_fields_and_values, other.comparison_fields_and_values
+        ):
+            if name != other_name:
+                comparison_fields = [n for n, v in self.comparison_fields_and_values]
+                other_comparison_fields = [n for n, v in other.comparison_fields_and_values]
+                raise ValueError(
+                    f"comparison field names do not match: {comparison_fields} != {other_comparison_fields}"
+                )
+            if value != other_value:
+                return value < other_value
+        return False
+
 
 T = TypeVar("T", covariant=False, bound="Annotation")
 

tests/core/test_document.py

@@ -266,6 +266,44 @@ class BinaryRelationWithEvidence(Annotation):
     _test_annotation_reconstruction(relation, annotation_store=annotation_store)
 
 
+def test_annotation_sort():
+    @dataclasses.dataclass(eq=True, frozen=True)
+    class Dummy(Annotation):
+        a: int
+        b: Optional[int] = None
+        c: int = 0
+
+    dummy1 = Dummy(a=1, c=2)
+    dummy2 = Dummy(a=1, c=3)
+    dummy3 = Dummy(a=2, c=1)
+    dummy4 = Dummy(a=2, c=2)
+
+    assert sorted([dummy1, dummy2, dummy3, dummy4]) == [dummy1, dummy2, dummy3, dummy4]
+
+    @dataclasses.dataclass(eq=True, frozen=True)
+    class DummyWithNestedAnnotation(Annotation):
+        a: int
+        n: Dummy
+
+    dummy_nested1 = DummyWithNestedAnnotation(a=1, n=Dummy(a=1, c=2))
+    dummy_nested2 = DummyWithNestedAnnotation(a=2, n=Dummy(a=2, c=3))
+    dummy_nested3 = DummyWithNestedAnnotation(a=2, n=Dummy(a=1, c=4))
+    dummy_nested4 = DummyWithNestedAnnotation(a=1, n=Dummy(a=2, c=2))
+
+    assert sorted([dummy_nested1, dummy_nested2, dummy_nested3, dummy_nested4]) == [
+        dummy_nested1,
+        dummy_nested4,
+        dummy_nested3,
+        dummy_nested2,
+    ]
+
+    with pytest.raises(ValueError) as excinfo:
+        sorted([dummy1, dummy_nested1])
+    assert (
+        str(excinfo.value) == "comparison field names do not match: ['a', 'n'] != ['a', 'b', 'c']"
+    )
+
+
 def test_annotation_is_attached():
     @dataclasses.dataclass
     class MyDocument(Document):