Document.add_all_annotations_from_other returns a mapping (#418)

* Document.add_all_annotations_from_other returns a mapping from original annotations to new ones

* make mypy happy

* improve tests

Author: ArneBinder

src/pytorch_ie/core/document.py

@@ -771,7 +771,7 @@ def add_all_annotations_from_other(
         process_predictions: bool = True,
         strict: bool = True,
         verbose: bool = True,
-    ) -> Dict[str, List[Annotation]]:
+    ) -> Dict[str, Dict[Annotation, Annotation]]:
         """Adds all annotations from another document to this document. It allows to blacklist annotations
         and also to override annotations. It returns the original annotations for which a new annotation was
         added to the current document.
@@ -854,7 +854,7 @@ class TokenBasedDocumentWithEntitiesRelationsAndRelationAttributes(TokenBasedDoc
                 ```
         """
         removed_annotations = defaultdict(set, removed_annotations or dict())
-        added_annotations = defaultdict(list)
+        added_annotations: Dict[str, Dict[Annotation, Annotation]] = defaultdict(dict)
 
         annotation_store: Dict[str, Dict[int, Annotation]] = defaultdict(dict)
         named_annotation_fields = {field.name: field for field in self.annotation_fields()}
@@ -897,7 +897,7 @@ class TokenBasedDocumentWithEntitiesRelationsAndRelationAttributes(TokenBasedDoc
                         if ann._id != new_ann._id:
                             annotation_store[field_name][ann._id] = new_ann
                         self[field_name].append(new_ann)
-                        added_annotations[field_name].append(ann)
+                        added_annotations[field_name][ann] = new_ann
                     else:
                         if strict:
                             raise ValueError(
@@ -922,7 +922,7 @@ class TokenBasedDocumentWithEntitiesRelationsAndRelationAttributes(TokenBasedDoc
                             if ann._id != new_ann._id:
                                 annotation_store[field_name][ann._id] = new_ann
                             self[field_name].predictions.append(new_ann)
-                            added_annotations[field_name].append(ann)
+                            added_annotations[field_name][ann] = new_ann
                         else:
                             if strict:
                                 raise ValueError(

tests/test_document.py

@@ -662,10 +662,13 @@ def test_document_extend_from_other_full_copy(text_document):
         "relation_attributes",
         "labels",
     }
-    for layer_name, annotation_set in added_annotations.items():
-        assert len(annotation_set) > 0
+    for layer_name, annotation_mapping in added_annotations.items():
+        assert len(annotation_mapping) > 0
         available_annotations = text_document[layer_name]
-        assert annotation_set == list(available_annotations)
+        assert set(annotation_mapping) == set(available_annotations)
+        assert len(annotation_mapping) == 1
+        # since we have only one annotation, we can construct the expected mapping
+        assert annotation_mapping == {available_annotations[0]: doc_new[layer_name][0]}
 
 
 def test_document_extend_from_other_wrong_override_annotation_mapping(text_document):
@@ -705,12 +708,19 @@ class TestDocument2(TokenBasedDocument):
     added_annotations = token_document.add_all_annotations_from_other(
         text_document, override_annotations=annotation_mapping
     )
+    added_annotation_sets = {k: set(v) for k, v in added_annotations.items()}
     # check that the added annotations are as expected (the entity annotations are already there)
-    assert added_annotations == {
-        "relations": list(text_document.relations),
-        "relation_attributes": list(text_document.relation_attributes),
-        "labels": list(text_document.labels),
+    assert added_annotation_sets == {
+        "relations": set(text_document.relations),
+        "relation_attributes": set(text_document.relation_attributes),
+        "labels": set(text_document.labels),
     }
+    for layer_name, annotation_mapping in added_annotations.items():
+        text_annotations = text_document[layer_name]
+        token_annotations = token_document[layer_name]
+        assert len(annotation_mapping) == len(text_annotations) == len(token_annotations) == 1
+        # since we have only one annotation, we can construct the expected mapping
+        assert annotation_mapping == {text_annotations[0]: token_annotations[0]}
 
     assert (
         len(token_document.entities1)
@@ -740,11 +750,15 @@ def test_document_extend_from_other_remove(text_document):
         removed_annotations={"entities1": {text_document.entities1[0]._id}},
         strict=False,
     )
-
+    added_annotation_sets = {k: set(v) for k, v in added_annotations.items()}
     # the only entity in entities1 is removed and since the relation has it as head, the relation is removed as well
+    assert added_annotation_sets == {
+        "entities2": set(text_document.entities2),
+        "labels": set(text_document.labels),
+    }
     assert added_annotations == {
-        "entities2": list(text_document.entities2),
-        "labels": list(text_document.labels),
+        "entities2": {text_document.entities2[0]: doc_new.entities2[0]},
+        "labels": {text_document.labels[0]: doc_new.labels[0]},
     }
 
     assert len(doc_new.entities1) == 0