1 files changed, 75 insertions, 0 deletions
diff --git a/coverage/annotate.py b/coverage/annotate.py
new file mode 100644
index 0000000..2d7d04a
--- /dev/null
+++ b/coverage/annotate.py
@@ -0,0 +1,75 @@
+"""Source file annotation for coverage.py"""
+
+import os, re
+
+from coverage.codeunit import code_unit_factory
+
+class AnnotateReporter:
+    def __init__(self, coverage, ignore_errors=False):
+        self.coverage = coverage
+        self.ignore_errors = ignore_errors
+        
+        self.directory = None
+        
+    blank_re = re.compile(r"\s*(#|$)")
+    else_re = re.compile(r"\s*else\s*:\s*(#|$)")
+
+    def report(self, morfs, directory=None, omit_prefixes=None):
+        morfs = morfs or self.coverage.data.executed_files()
+        code_units = code_unit_factory(morfs, self.coverage.file_locator, omit_prefixes)
+        self.directory = directory
+        for cu in code_units:
+            try:
+                filename, statements, excluded, missing, _ = self.coverage.analyze(cu)
+                self.annotate_file(filename, statements, excluded, missing)
+            except KeyboardInterrupt:
+                raise
+            except:
+                if not self.ignore_errors:
+                    raise
+                
+    def annotate_file(self, filename, statements, excluded, missing):
+        source = open(filename, 'r')
+        if self.directory:
+            dest_file = os.path.join(self.directory,
+                                     os.path.basename(filename)
+                                     + ',cover')
+        else:
+            dest_file = filename + ',cover'
+        dest = open(dest_file, 'w')
+        lineno = 0
+        i = 0
+        j = 0
+        covered = True
+        while True:
+            line = source.readline()
+            if line == '':
+                break
+            lineno = lineno + 1
+            while i < len(statements) and statements[i] < lineno:
+                i = i + 1
+            while j < len(missing) and missing[j] < lineno:
+                j = j + 1
+            if i < len(statements) and statements[i] == lineno:
+                covered = j >= len(missing) or missing[j] > lineno
+            if self.blank_re.match(line):
+                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):
+                    dest.write('> ')
+                elif statements[i] == missing[j]:
+                    dest.write('! ')
+                else:
+                    dest.write('> ')
+            elif lineno in excluded:
+                dest.write('- ')
+            elif covered:
+                dest.write('> ')
+            else:
+                dest.write('! ')
+            dest.write(line)
+        source.close()
+        dest.close()