Add a check for overlapping annotations

Signed-off-by: Konstantin Slavnov <[email protected]>

lookout/style/format/annotations/annotated_data.py

@@ -7,12 +7,12 @@
     LanguageAnnotation, PathAnnotation, UASTAnnotation
 
 
-class NoIntersection(Exception):
+class NoAnnotation(Exception):
     """
-    Raised by `AnnotationManager.find_intersect()` if there is no intersection of provided \
-    intervals.
+    Raised by `AnnotationManager` methods if there is no Annotation found.
 
-    See documentation about `AnnotationManager.find_intersect()` for more information.
+    See documentation about `AnnotationManager.find_overlapping_annotation()` or
+    `AnnotationManager.find_covering_annotation()` for more information.
     """
 
 
@@ -98,7 +98,10 @@ def __getitem__(self, item: Union[int, slice, Tuple[int, int]]) -> str:
 
     def count(self, annotation_type: Type[Annotation]):
         """Count the number of annotations of a specific type."""
-        return len(self._type_to_annotations[annotation_type])
+        try:
+            return len(self._type_to_annotations[annotation_type])
+        except KeyError:
+            return 0
 
     def add(self, *annotations: Annotation) -> None:
         """
@@ -117,11 +120,26 @@ def _add(self, annotation: Annotation) -> None:
                 raise ValueError("Global annotation %s already exists" % annotation)
             self._global_annotations[annotation_type] = annotation
         else:
-            # TODO(zurk): Add a check that there is no overlapping annotations of one type.
             if annotation.span not in self._span_to_annotations:
                 self._span_to_annotations[annotation.span] = {}
             if annotation_type not in self._type_to_annotations:
                 self._type_to_annotations[annotation_type] = SortedDict()
+            if self._type_to_annotations[annotation_type]:
+                right_span_index = self._type_to_annotations[annotation_type].bisect_left(
+                    annotation.span)
+                if right_span_index < len(self._type_to_annotations[annotation_type]):
+                    right_span, right_annotation = \
+                        self._type_to_annotations[annotation_type].peekitem(right_span_index)
+                    if self._check_spans_overlap(*right_span, *annotation.span):
+                        raise ValueError(
+                            "Trying to add Annotation %s which is overlaps with existing %s" % (
+                                right_annotation, annotation))
+                left_span_index = max(0, right_span_index - 1)
+                left_span, left_annotation = self._type_to_annotations[annotation_type].peekitem(
+                    left_span_index)
+                if self._check_spans_overlap(*left_span, *annotation.span):
+                    raise ValueError("Trying to add Annotation %s which is overlaps with existing"
+                                     "%s" % (left_annotation, annotation))
             self._span_to_annotations[annotation.span][annotation_type] = annotation
             self._type_to_annotations[annotation_type][annotation.span] = annotation
 
@@ -173,9 +191,9 @@ def iter_annotations(self, annotation_type: Type[Annotation],
             annotations = dict()
             for missing_type in missing_types:
                 try:
-                    annotations[missing_type] = self.find_overlapping_span(missing_type,
-                                                                           *annotation.span)
-                except NoIntersection:
+                    annotations[missing_type] = self.find_covering_annotation(missing_type,
+                                                                              *annotation.span)
+                except NoAnnotation:
                     pass
             annotations.update({type: same_span_annotations[type] for type in common_types})
             yield AnnotationsSpan(*annotation.span, annotations)
@@ -201,30 +219,58 @@ def iter_annotation(self, annotation_type: Type[Annotation], *,
         for value in self._type_to_annotations[annotation_type].values()[search_from:]:
             yield value
 
-    def find_overlapping_span(self, annotation_type: Type[Annotation],
-                              start: int, stop: int) -> Annotation:
+    def find_overlapping_annotation(self, annotation_type: Type[Annotation],
+                                    start: int, stop: int) -> Annotation:
         """
         Find an annotation of the given type that intersects the interval [start, stop).
 
         :param annotation_type: Annotation type to look for.
         :param start: Start of the search interval.
         :param stop: End of the search interval. Stop point itself is excluded.
-        :raise NoIntersection: There is no such annotation that overlaps with the given interval.
+        :raise NoAnnotation: There is no such annotation that overlaps with the given interval.
         :return: `Annotation` of the requested type.
         """
         try:
             annotation_layer = self._type_to_annotations[annotation_type]
         except KeyError:
-            raise NoIntersection("There is no annotation layer %s" % annotation_type)
+            raise NoAnnotation("There is no annotation layer %s" % annotation_type)
         check_span(start, stop)
         search_start = max(0, annotation_layer.bisect_left((start, start)) - 1)
         search_stop = annotation_layer.bisect_right((stop, stop))
         for span in annotation_layer.islice(search_start, search_stop):
             if self._check_spans_overlap(start, stop, *span):
                 # assuming that there is only one such annotation
                 return annotation_layer[span]
-        raise NoIntersection("There is no annotation %s from %d to %d" % (annotation_type, start,
-                                                                          stop))
+        raise NoAnnotation("There is no annotation %s from %d to %d" % (
+            annotation_type, start, stop))
+
+    def find_covering_annotation(self, annotation_type: Type[Annotation],
+                                 start: int, stop: int) -> Annotation:
+        """
+        Find an annotation of the given type that fully covers the interval [start, stop).
+
+        :param annotation_type: Annotation type to look for.
+        :param start: Start of the search interval.
+        :param stop: End of the search interval. Stop point itself is excluded.
+        :raise NoAnnotation: There is no such annotation that overlaps with the given interval.
+        :return: `Annotation` of the requested type.
+        """
+        try:
+            annotation_layer = self._type_to_annotations[annotation_type]
+        except KeyError:
+            raise NoAnnotation("There is no annotation layer %s" % annotation_type)
+        check_span(start, stop)
+        search_start = max(0, annotation_layer.bisect_left((start, start)) - 1)
+        search_stop = annotation_layer.bisect_right((stop, stop))
+        for span_start, span_stop in annotation_layer.islice(search_start, search_stop):
+            if start == stop:
+                if self._check_spans_overlap(start, stop, span_start, span_stop):
+                    return annotation_layer[(span_start, span_stop)]
+            elif span_start <= start and stop <= span_stop:
+                # assuming that there is only one such annotation because they cannot overlap
+                return annotation_layer[(span_start, span_stop)]
+        raise NoAnnotation("There is no annotation %s that contains i%d to %d" % (
+            annotation_type, start, stop))
 
     @classmethod
     def _check_spans_overlap(cls, start1: int, stop1: int, start2: int, stop2: int) -> bool: