Incremental syntax highlight for IRB source lines

Closes: https://github.com/ruby/ruby/pull/2202

NEWS

@@ -115,8 +115,8 @@ ERB::
 
 IRB::
 
-  * Introduce syntax highlight inspired by pry.gem to inspect output for some
-    core-class objects and binding.irb source lines if $TERM is set and not dumb.
+  * Introduce syntax highlight inspired by pry.gem to binding.irb source lines,
+    REPL input, and inspect output of some core-class objects.
 
 Net::IMAP::
 

lib/irb/color.rb

@@ -77,13 +77,21 @@ module IRB # :nodoc:
         return code unless colorable?
 
         colored = +''
+        length = 0
         Ripper.lex(code).each do |(_line, _col), token, str, expr|
           if seq = dispatch_seq(token, expr, str)
-            colored << "#{seq.map { |s| "\e[#{s}m" }.join('')}#{str}#{clear}"
+            str.each_line do |line|
+              colored << "#{seq.map { |s| "\e[#{s}m" }.join('')}#{line}#{clear}"
+            end
           else
             colored << str
           end
+          length += str.length
         end
+
+        # give up colorizing incomplete Ripper tokens
+        return code if length != code.length
+
         colored
       end
 

lib/irb/input-method.rb

@@ -222,6 +222,10 @@ module IRB
       end
       Reline.completion_append_character = nil
       Reline.completion_proc = IRB::InputCompletor::CompletionProc
+      Reline.output_modifier_proc = proc do |output|
+        next unless IRB::Color.colorable?
+        IRB::Color.colorize_code(output)
+      end
       Reline.dig_perfect_match_proc = IRB::InputCompletor::PerfectMatchedProc
     end
 

lib/reline.rb

@@ -142,6 +142,15 @@ module Reline
     @@completion_proc = p
   end
 
+  @@output_modifier_proc = nil
+  def self.output_modifier_proc
+    @@output_modifier_proc
+  end
+  def self.output_modifier_proc=(p)
+    raise ArgumentError unless p.is_a?(Proc)
+    @@output_modifier_proc = p
+  end
+
   @@pre_input_hook = nil
   def self.pre_input_hook
     @@pre_input_hook
@@ -297,6 +306,7 @@ module Reline
     end
     @@line_editor.output = @@output
     @@line_editor.completion_proc = @@completion_proc
+    @@line_editor.output_modifier_proc = @@output_modifier_proc
     @@line_editor.dig_perfect_match_proc = @@dig_perfect_match_proc
     @@line_editor.pre_input_hook = @@pre_input_hook
     @@line_editor.retrieve_completion_block = method(:retrieve_completion_block)

lib/reline/line_editor.rb

@@ -10,6 +10,7 @@ class Reline::LineEditor
   attr_reader :byte_pointer
   attr_accessor :confirm_multiline_termination_proc
   attr_accessor :completion_proc
+  attr_accessor :output_modifier_proc
   attr_accessor :pre_input_hook
   attr_accessor :dig_perfect_match_proc
   attr_writer :retrieve_completion_block
@@ -163,16 +164,16 @@ class Reline::LineEditor
     lines = [String.new(encoding: @encoding)]
     height = 1
     width = 0
-    rest_prompt = prompt.encode(Encoding::UTF_8)
+    rest = "#{prompt}#{str}".encode(Encoding::UTF_8)
     loop do
-      break if rest_prompt.empty?
-      if rest_prompt =~ /^#{CSI_REGEXP}/
+      break if rest.empty?
+      if rest =~ /\A#{CSI_REGEXP}/
         lines.last << $&
-        rest_prompt = $'
+        rest = $'
       else
-        gcs = rest_prompt.grapheme_clusters
+        gcs = rest.grapheme_clusters
         gc = gcs.first
-        rest_prompt = gcs[1..-1].join
+        rest = gcs[1..-1].join
         mbchar_width = Reline::Unicode.get_mbchar_width(gc)
         width += mbchar_width
         if width > max_width
@@ -184,19 +185,6 @@ class Reline::LineEditor
         lines.last << gc
       end
     end
-    lines << :split
-    lines << String.new(encoding: @encoding)
-    str.encode(Encoding::UTF_8).grapheme_clusters.each do |gc|
-      mbchar_width = Reline::Unicode.get_mbchar_width(gc)
-      width += mbchar_width
-      if width > max_width
-        width = mbchar_width
-        lines << nil
-        lines << String.new(encoding: @encoding)
-        height += 1
-      end
-      lines.last << gc
-    end
     # The cursor moves to next line in first
     lines << String.new(encoding: @encoding) if width == max_width
     [lines, height]
@@ -290,8 +278,7 @@ class Reline::LineEditor
       Reline::IOGate.clear_screen
       @cleared = false
       back = 0
-      @buffer_of_lines.each_with_index do |line, index|
-        line = @line if index == @line_index
+      modify_lines(whole_lines).each_with_index do |line, index|
         height = render_partial(prompt, prompt_width, line, false)
         if index < (@buffer_of_lines.size - 1)
           move_cursor_down(height)
@@ -305,7 +292,6 @@ class Reline::LineEditor
     end
     # FIXME: end of logical line sometimes breaks
     if @previous_line_index
-      previous_line = @line
       all_height = @buffer_of_lines.inject(0) { |result, line|
         result + calculate_height_by_width(@prompt_width + calculate_width(line))
       }
@@ -316,8 +302,7 @@ class Reline::LineEditor
         scroll_down(diff)
         move_cursor_up(@highest_in_all - 1)
         back = 0
-        @buffer_of_lines.each_with_index do |line, index|
-          line = @line if index == @previous_line_index
+        modify_lines(whole_lines(index: @previous_line_index, line: @line)).each_with_index do |line, index|
           height = render_partial(prompt, prompt_width, line, false)
           if index < (@buffer_of_lines.size - 1)
             move_cursor_down(height)
@@ -326,6 +311,7 @@ class Reline::LineEditor
         end
         move_cursor_up(back)
       else
+        previous_line = modify_lines(whole_lines(index: @previous_line_index, line: @line))[@previous_line_index]
         render_partial(prompt, prompt_width, previous_line)
         move_cursor_up(@first_line_started_from + @started_from)
       end
@@ -364,7 +350,7 @@ class Reline::LineEditor
         end
         move_cursor_up(@highest_in_all - 1)
       end
-      @buffer_of_lines.each_with_index do |line, index|
+      modify_lines(@buffer_of_lines).each_with_index do |line, index|
         render_partial(prompt, prompt_width, line, false)
         if index < (@buffer_of_lines.size - 1)
           move_cursor_down(1)
@@ -384,10 +370,11 @@ class Reline::LineEditor
       move_cursor_down(@first_line_started_from)
       @rerender_all = false
     end
+    line = modify_lines(whole_lines)[@line_index]
     if !@is_multiline
-      render_partial(prompt, prompt_width, @line)
+      render_partial(prompt, prompt_width, line)
     elsif !finished?
-      render_partial(prompt, prompt_width, @line)
+      render_partial(prompt, prompt_width, line)
     else
       scroll_down(1) unless whole_lines.last.empty?
       Reline::IOGate.move_cursor_column(0)
@@ -408,24 +395,15 @@ class Reline::LineEditor
       move_cursor_up(@started_from)
       @started_from = calculate_height_by_width(prompt_width + @cursor) - 1
     end
-    is_prompt = true
     Reline::IOGate.move_cursor_column(0)
     visual_lines.each do |line|
-      if line == :split
-        is_prompt = false
-        next
-      end
       if line.nil?
         Reline::IOGate.erase_after_cursor
         move_cursor_down(1)
         Reline::IOGate.move_cursor_column(0)
         next
       end
-      if is_prompt
-        @output.print line
-      else
-        escaped_print line
-      end
+      @output.print line
       if @first_prompt
         @first_prompt = false
         @pre_input_hook&.call
@@ -441,6 +419,16 @@ class Reline::LineEditor
     height
   end
 
+  private def modify_lines(before)
+    return before if before.nil? || before.empty?
+
+    if after = @output_modifier_proc&.call("#{before.join("\n")}\n")
+      after.lines(chomp: true)
+    else
+      before
+    end
+  end
+
   def editing_mode
     @config.editing_mode
   end
@@ -768,9 +756,9 @@ class Reline::LineEditor
     @cursor_max = calculate_width(@line)
   end
 
-  def whole_lines
+  def whole_lines(index: @line_index, line: @line)
     temp_lines = @buffer_of_lines.dup
-    temp_lines[@line_index] = @line
+    temp_lines[index] = line
     temp_lines
   end
 

test/irb/test_color.rb

@@ -32,6 +32,9 @@ module TestIRB
         '"##@var]"' => "#{RED}\"#{CLEAR}#{RED}##{CLEAR}#{RED}##{CLEAR}@var#{RED}]#{CLEAR}#{RED}\"#{CLEAR}",
         '"foo#{a} #{b}"' => "#{RED}\"#{CLEAR}#{RED}foo#{CLEAR}#{RED}\#{#{CLEAR}a#{RED}}#{CLEAR}#{RED} #{CLEAR}#{RED}\#{#{CLEAR}b#{RED}}#{CLEAR}#{RED}\"#{CLEAR}",
         '/r#{e}g/' => "#{RED}#{BOLD}/#{CLEAR}#{RED}r#{CLEAR}#{RED}\#{#{CLEAR}e#{RED}}#{CLEAR}#{RED}g#{CLEAR}#{RED}#{BOLD}/#{CLEAR}",
+        "'a\nb'" => "#{RED}'#{CLEAR}#{RED}a\n#{CLEAR}#{RED}b#{CLEAR}#{RED}'#{CLEAR}",
+        "4.5.6" => "4.5.6",
+        "[1]]]" => "[1]]]",
       }.each do |code, result|
         assert_equal(result, with_term { IRB::Color.colorize_code(code) }, "Case: colorize_code(#{code.dump})")
       end