From bfb0cfa4fd65089094fd6f552e72689d642b9788 Mon Sep 17 00:00:00 2001
From: Jules Villard <jul@fb.com>
Date: Wed, 9 Nov 2016 07:36:44 -0800
Subject: [PATCH] [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 | 9 ++++++++-
 infer/lib/python/inferlib/source.py   | 7 ++++++-
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/infer/lib/python/inferlib/colorize.py b/infer/lib/python/inferlib/colorize.py
index 4af9c7951..0c67ef5e3 100644
--- a/infer/lib/python/inferlib/colorize.py
+++ b/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):
diff --git a/infer/lib/python/inferlib/source.py b/infer/lib/python/inferlib/source.py
index c8b9d6301..2466855e3 100644
--- a/infer/lib/python/inferlib/source.py
+++ b/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'):