[python] work around bug in pygments to display the right source excerpt

Summary:
The syntax highlighting swallows initial '\n' characters in the string, which
caused bugs where infer would report the correct line but not the correct
source excerpt.

Reviewed By: jeremydubreil

Differential Revision: D4153083

fbshipit-source-id: 8b1d211

infer/lib/python/inferlib/colorize.py

@@ -66,7 +66,14 @@ def syntax_highlighting(source_name, mode, s):
         if not sys.stdout.isatty():
             return s
         formatter = pygments.formatters.TerminalFormatter()
-    return pygments.highlight(s, lexer, formatter)
+    # there's a bug in pygments.highlight() where it will remove all starting
+    # newline characters, so we have to add them back!
+    initial_newlines = ''
+    i = 0
+    while (i < len(s) and s[i] == '\n'):
+        initial_newlines += '\n'
+        i += 1
+    return initial_newlines + pygments.highlight(s, lexer, formatter)
 
 
 def color(s, color, mode):

infer/lib/python/inferlib/source.py

@@ -71,13 +71,18 @@ def build_source_context(source_name, mode, report_line):
                      encoding=config.CODESET, errors="replace") as source_file:
         # avoid going past the end of the file
         for line in source_file:
+            last_line = line_number
             if start_line <= line_number <= end_line:
                 excerpt += line
+            elif line_number > end_line:
+                # OPTIM: no need to read past the last line of the excerpt
+                break
             line_number += 1
     excerpt = colorize.syntax_highlighting(source_name, mode, excerpt)
 
+
     # number lines and add caret at the right position
-    n_length = len(str(end_line))
+    n_length = len(str(last_line))
     s = ''
     line_number = start_line
     for line in excerpt.split('\n'):