Apply syntax highlighting when expanding diff plus some refactor. #3945

1 files changed, 43 insertions, 12 deletions

diff --git a/lib/gitlab/diff/highlight.rb b/lib/gitlab/diff/highlight.rb
index c780ea21775..7f340de65cc 100644
--- a/lib/gitlab/diff/highlight.rb
+++ b/lib/gitlab/diff/highlight.rb
@@ -1,31 +1,62 @@
 module Gitlab
   module Diff
     class Highlight
-      def self.process_diff_lines(file_name, diff_lines)
-        processor = new(file_name, diff_lines)
+      # Apply syntax highlight to provided source code
+      #
+      # file_name - The file name related to the code.
+      # lines     - It can be an Array of Gitlab::Diff::Line objects or simple Strings.
+      #             When passing Strings you need to provide the required 'end of lines'
+      #             chars ("\n") for each String given that we don't append them automatically.
+      #
+      # Returns an Array with the processed items.
+      def self.process_diff_lines(file_name, lines)
+        processor = new(file_name, lines)
         processor.highlight
       end
 
-      def initialize(file_name, diff_lines)
-        text_lines          = diff_lines.map(&:text)
-        @file_name          = file_name
-        @diff_lines         = diff_lines
-        @diff_line_prefixes = text_lines.map { |line| line.sub!(/\A((\+|\-)\s*)/, '');$1 }
-        @raw_lines          = text_lines.join("\n")
+      def initialize(file_name, lines)
+        @file_name  = file_name
+        @lines      = lines
       end
 
       def highlight
-        @code = unescape_html(@raw_lines)
+        return [] if @lines.empty?
+
+        extract_line_prefixes
+
+        @code             = unescape_html(raw_content)
         @highlighted_code = formatter.format(lexer.lex(@code))
 
-        update_diff_lines
+        is_diff_line? ? update_diff_lines : @highlighted_code.lines
       end
 
       private
 
+      def is_diff_line?
+        @lines.first.is_a?(Gitlab::Diff::Line)
+      end
+
+      def text_lines
+        @text_lines ||= (is_diff_line? ? @lines.map(&:text) : @lines)
+      end
+
+      def raw_content
+        @raw_content ||= text_lines.join(is_diff_line? ? "\n" : nil)
+      end
+
+      def extract_line_prefixes
+        @diff_line_prefixes ||= begin
+          if is_diff_line?
+            text_lines.map { |line| line.sub!(/\A((\+|\-)\s*)/, '');$1 }
+          else
+            []
+          end
+        end
+      end
+
       def update_diff_lines
         @highlighted_code.lines.each_with_index do |line, i|
-          diff_line = @diff_lines[i]
+          diff_line = @lines[i]
 
           # ignore highlighting for "match" lines
           next if diff_line.type == 'match'
@@ -33,7 +64,7 @@ module Gitlab
           diff_line.text = "#{@diff_line_prefixes[i]}#{line}"
         end
 
-        @diff_lines
+        @lines
       end
 
       def lexer