test: add a annotate test, and simplify the logic

Author: nedbat

coverage/annotate.py

@@ -74,8 +74,7 @@ def annotate_file(self, fr, analysis):
             dest_file = fr.filename + ",cover"
 
         with open(dest_file, 'w', encoding='utf-8') as dest:
-            i = 0
-            j = 0
+            i = j = 0
             covered = True
             source = fr.source()
             for lineno, line in enumerate(source.splitlines(True), start=1):
@@ -89,9 +88,7 @@ def annotate_file(self, fr, analysis):
                     dest.write('  ')
                 elif self.else_re.match(line):
                     # Special logic for lines containing only 'else:'.
-                    if i >= len(statements) and j >= len(missing):
-                        dest.write('! ')
-                    elif i >= len(statements) or j >= len(missing):
+                    if j >= len(missing):
                         dest.write('> ')
                     elif statements[i] == missing[j]:
                         dest.write('! ')

tests/gold/annotate/mae/mae.py,cover

@@ -0,0 +1,10 @@
+> def f(x):
+>     if x == 1:
+>         print("1")
+>     else:
+>         print("2")
+  
+> if f(1):
+!     print("nope")
+> if f(2):
+!     print("nope")

tests/test_annotate.py

@@ -105,3 +105,23 @@ def h(x):
         self.start_import_stop(cov, "white")
         cov.annotate()
         compare(gold_path("annotate/annotate"), ".", "*,cover")
+
+    def test_missing_after_else(self):
+        self.make_file("mae.py", """\
+            def f(x):
+                if x == 1:
+                    print("1")
+                else:
+                    print("2")
+
+            if f(1):
+                print("nope")
+            if f(2):
+                print("nope")
+            """)
+
+        cov = coverage.Coverage()
+        self.start_import_stop(cov, "mae")
+        cov.annotate()
+        assert self.stdout() == "1\n2\n"
+        compare(gold_path("annotate/mae"), ".", "*,cover")