[ruby/irb] Fix pager preview with escape sequence and newlines

(https://github.com/ruby/irb/pull/1069)

https://github.com/ruby/irb/commit/a139562a07

2 files changed, 19 insertions, 8 deletions

diff --git a/lib/irb/pager.rb b/lib/irb/pager.rb
index 65303e5ac1..16ff30cf89 100644
--- a/lib/irb/pager.rb
+++ b/lib/irb/pager.rb
@@ -152,11 +152,13 @@ module IRB
       end
 
       def puts(text = '')
+        text = text.to_s unless text.is_a?(String)
         write(text)
         write("\n") unless text.end_with?("\n")
       end
 
       def write(text)
+        text = text.to_s unless text.is_a?(String)
         @string << text
         if @multipage
           if @delay_until && Time.now > @delay_until
@@ -171,23 +173,24 @@ module IRB
           text = text[0, overflow_size]
           overflow = true
         end
-
         @buffer << text
-        @col += Reline::Unicode.calculate_width(text)
+        @col += Reline::Unicode.calculate_width(text, true)
         if text.include?("\n") || @col >= @width
           @buffer.lines.each do |line|
             wrapped_lines = Reline::Unicode.split_by_width(line.chomp, @width).first.compact
             wrapped_lines.pop if wrapped_lines.last == ''
             @lines.concat(wrapped_lines)
-            if @lines.empty?
-              @lines << "\n"
-            elsif line.end_with?("\n")
-              @lines[-1] += "\n"
+            if line.end_with?("\n")
+              if @lines.empty? || @lines.last.end_with?("\n")
+                @lines << "\n"
+              else
+                @lines[-1] += "\n"
+              end
             end
           end
           @buffer.clear
           @buffer << @lines.pop unless @lines.last.end_with?("\n")
-          @col = Reline::Unicode.calculate_width(@buffer)
+          @col = Reline::Unicode.calculate_width(@buffer, true)
         end
         if overflow || @lines.size > @height || (@lines.size == @height && @col > 0)
           @first_page_lines = @lines.take(@height)
diff --git a/test/irb/test_pager.rb b/test/irb/test_pager.rb
index 5842519e62..0fad94da30 100644
--- a/test/irb/test_pager.rb
+++ b/test/irb/test_pager.rb
@@ -7,13 +7,21 @@ module TestIRB
   class PagerTest < TestCase
     def test_take_first_page
       assert_equal ['a' * 40, true], IRB::Pager.take_first_page(10, 4) {|io| io.puts 'a' * 41; raise 'should not reach here' }
+      assert_equal ["a\nb\na\nb\n", true], IRB::Pager.take_first_page(10, 4) {|io| 10.times { io.puts "a\nb\n" } }
+      assert_equal ["a\n\n\na\n", true], IRB::Pager.take_first_page(10, 4) {|io| 10.times { io.puts "a\n\n\n" } }
+      assert_equal ["11\n" * 4, true], IRB::Pager.take_first_page(10, 4) {|io| 10.times { io.write 1; io.puts 1 } }
+      assert_equal ["\n" * 4, true], IRB::Pager.take_first_page(10, 4) {|io| 10.times { io.write nil; io.puts nil } }
       assert_equal ['a' * 39, false], IRB::Pager.take_first_page(10, 4) {|io| io.write 'a' * 39 }
       assert_equal ['a' * 39 + 'b', false], IRB::Pager.take_first_page(10, 4) {|io| io.write 'a' * 39 + 'b' }
       assert_equal ['a' * 39 + 'b', true], IRB::Pager.take_first_page(10, 4) {|io| io.write 'a' * 39 + 'bc' }
       assert_equal ["a\nb\nc\nd\n", false], IRB::Pager.take_first_page(10, 4) {|io| io.write "a\nb\nc\nd\n" }
       assert_equal ["a\nb\nc\nd\n", true], IRB::Pager.take_first_page(10, 4) {|io| io.write "a\nb\nc\nd\ne" }
       assert_equal ['a' * 15 + "\n" + 'b' * 20, true], IRB::Pager.take_first_page(10, 4) {|io| io.puts 'a' * 15; io.puts 'b' * 30 }
-      assert_equal ["\e[31mA\e[0m" * 10 + 'x' * 30, true], IRB::Pager.take_first_page(10, 4) {|io| io.puts "\e[31mA\e[0m" * 10 + 'x' * 31; }
+      assert_equal ["\e[31mA\e[0m" * 10 + 'x' * 30, true], IRB::Pager.take_first_page(10, 4) {|io| io.puts "\e[31mA\e[0m" * 10 + 'x' * 31 }
+      text, overflow = IRB::Pager.take_first_page(10, 4) {|io| 41.times { io.write "\e[31mA\e[0m" } }
+      assert_equal ['A' * 40, true], [text.gsub(/\e\[\d+m/, ''), overflow]
+      text, overflow = IRB::Pager.take_first_page(10, 4) {|io| 41.times { io.write "\e[31mAAA\e[0m" } }
+      assert_equal ['A' * 40, true], [text.gsub(/\e\[\d+m/, ''), overflow]
     end
   end