Quickly show inspect preview even if pretty_print takes too much time

lib/irb.rb

@@ -517,40 +517,36 @@ def signal_status(status)
     end
 
     def output_value(omit = false) # :nodoc:
-      str = @context.inspect_last_value
-      multiline_p = str.include?("\n")
-      if omit
-        winwidth = @context.io.winsize.last
-        if multiline_p
-          first_line = str.split("\n").first
-          result = @context.newline_before_multiline_output? ? (@context.return_format % first_line) : first_line
-          output_width = Reline::Unicode.calculate_width(result, true)
-          diff_size = output_width - Reline::Unicode.calculate_width(first_line, true)
-          if diff_size.positive? and output_width > winwidth
-            lines, _ = Reline::Unicode.split_by_width(first_line, winwidth - diff_size - 3)
-            str = "%s..." % lines.first
-            str += "\e[0m" if Color.colorable?
-            multiline_p = false
-          else
-            str = str.gsub(/(\A.*?\n).*/m, "\\1...")
-            str += "\e[0m" if Color.colorable?
-          end
-        else
-          output_width = Reline::Unicode.calculate_width(@context.return_format % str, true)
-          diff_size = output_width - Reline::Unicode.calculate_width(str, true)
-          if diff_size.positive? and output_width > winwidth
-            lines, _ = Reline::Unicode.split_by_width(str, winwidth - diff_size - 3)
-            str = "%s..." % lines.first
-            str += "\e[0m" if Color.colorable?
-          end
-        end
+      unless @context.return_format.include?('%')
+        puts @context.return_format
+        return
       end
 
-      if multiline_p && @context.newline_before_multiline_output?
-        str = "\n" + str
+      winheight, winwidth = @context.io.winsize
+      if omit
+        content, overflow = Pager.take_first_page(winwidth, 1) do |out|
+          @context.inspect_last_value(out)
+        end
+        if overflow
+          content = "\n#{content}" if @context.newline_before_multiline_output?
+          content = "#{content}..."
+          content = "#{content}\e[0m" if Color.colorable?
+        end
+        puts format(@context.return_format, content.chomp)
+      elsif Pager.should_page? && @context.inspector_support_stream_output?
+        modifier_proc = ->(content, multipage) do
+          content = content.chomp
+          content = "\n#{content}" if @context.newline_before_multiline_output? && (multipage || content.include?("\n"))
+          format(@context.return_format, content)
+        end
+        Pager.page_with_preview(winwidth, winheight, modifier_proc) do |out|
+          @context.inspect_last_value(out)
+        end
+      else
+        content = @context.inspect_last_value.chomp
+        content = "\n#{content}" if @context.newline_before_multiline_output? && content.include?("\n")
+        Pager.page_content(format(@context.return_format, content), retain_content: true)
       end
-
-      Pager.page_content(format(@context.return_format, str), retain_content: true)
     end
 
     # Outputs the local variables to this current session, including #signal_status

lib/irb/context.rb

@@ -673,8 +673,12 @@ def colorize_input(input, complete:)
       end
     end
 
-    def inspect_last_value # :nodoc:
-      @inspect_method.inspect_value(@last_value)
+    def inspect_last_value(output = +'') # :nodoc:
+      @inspect_method.inspect_value(@last_value, output)
+    end
+
+    def inspector_support_stream_output?
+      @inspect_method.support_stream_output?
     end
 
     NOPRINTING_IVARS = ["@last_value"] # :nodoc:

lib/irb/inspector.rb

@@ -92,9 +92,13 @@ def init
       @init.call if @init
     end
 
+    def support_stream_output?
+      @inspect.arity == 2
+    end
+
     # Proc to call when the input is evaluated and output in irb.
-    def inspect_value(v)
-      @inspect.call(v)
+    def inspect_value(v, output)
+      @inspect.arity == 2 ? @inspect.call(v, output) : output << @inspect.call(v)
     rescue => e
       puts "An error occurred when inspecting the object: #{e.inspect}"
 
@@ -113,8 +117,8 @@ def inspect_value(v)
   Inspector.def_inspector([:p, :inspect]){|v|
     Color.colorize_code(v.inspect, colorable: Color.colorable? && Color.inspect_colorable?(v))
   }
-  Inspector.def_inspector([true, :pp, :pretty_inspect], proc{require_relative "color_printer"}){|v|
-    IRB::ColorPrinter.pp(v, +'').chomp
+  Inspector.def_inspector([true, :pp, :pretty_inspect], proc{require_relative "color_printer"}){|v, output|
+    IRB::ColorPrinter.pp(v, output)
   }
   Inspector.def_inspector([:yaml, :YAML], proc{require "yaml"}){|v|
     begin

lib/irb/pager.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require 'reline'
+
 module IRB
   # The implementation of this class is borrowed from RDoc's lib/rdoc/ri/driver.rb.
   # Please do NOT use this class directly outside of IRB.
@@ -47,12 +49,39 @@ def page(retain_content: false)
       rescue Errno::EPIPE
       end
 
-      private
-
       def should_page?
         IRB.conf[:USE_PAGER] && STDIN.tty? && (ENV.key?("TERM") && ENV["TERM"] != "dumb")
       end
 
+      def page_with_preview(width, height, modifier_proc)
+        out = PageOverflowIO.new(width, height) do |lines|
+          modified_output = modifier_proc.call(lines.join, true)
+          content, = take_first_page(width, height - 1) {|o| o.write modified_output }
+          content = content.chomp
+          content = "#{content}\e[0m" if Color.colorable?
+          $stdout.puts content
+        end
+        yield out
+        content = modifier_proc.call(out.string, out.multipage?)
+        if out.multipage?
+          page(retain_content: true) do |io|
+            io.puts content
+          end
+        else
+          $stdout.puts content
+        end
+      end
+
+      def take_first_page(width, height)
+        out = Pager::PageOverflowIO.new(width, height) do |lines|
+          return lines.join, true
+        end
+        yield out
+        [out.string, false]
+      end
+
+      private
+
       def content_exceeds_screen_height?(content)
         screen_height, screen_width = begin
           Reline.get_screen_size
@@ -62,10 +91,10 @@ def content_exceeds_screen_height?(content)
 
         pageable_height = screen_height - 3 # leave some space for previous and the current prompt
 
-        # If the content has more lines than the pageable height
-        content.lines.count > pageable_height ||
-          # Or if the content is a few long lines
-          pageable_height * screen_width < Reline::Unicode.calculate_width(content, true)
+        return true if content.lines.size > pageable_height
+
+        _, overflow = take_first_page(screen_width, pageable_height) {|out| out.write content }
+        overflow
       end
 
       def setup_pager(retain_content:)
@@ -96,5 +125,71 @@ def setup_pager(retain_content:)
         nil
       end
     end
+
+    # Writable IO that has page overflow callback
+    class PageOverflowIO
+      attr_reader :string
+
+      # Maximum size of a single cell in terminal
+      # Assumed worst case: "\e[1;3;4;9;38;2;255;128;128;48;2;128;128;255mA\e[0m"
+      # bold, italic, underline, crossed_out, RGB forgound, RGB background
+      MAX_CHAR_PER_CELL = 50
+
+      def initialize(width, height, &overflow_callback)
+        @lines = []
+        @width = width
+        @height = height
+        @buffer = +''
+        @overflow_callback = overflow_callback
+        @col = 0
+        @string = +''
+        @multipage = false
+      end
+
+      def puts(text = '')
+        write(text)
+        write("\n") unless text.end_with?("\n")
+      end
+
+      def write(text)
+        @string << text
+        return if @multipage
+
+        overflow_size = (@width * (@height - @lines.size) + @width - @col )* MAX_CHAR_PER_CELL
+        if text.size >= overflow_size
+          text = text[0, overflow_size]
+          overflow = true
+        end
+
+        @buffer << text
+        @col += Reline::Unicode.calculate_width(text)
+        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"
+            end
+          end
+          @buffer.clear
+          @buffer << @lines.pop unless @lines.last.end_with?("\n")
+          @col = Reline::Unicode.calculate_width(@buffer)
+        end
+        if overflow || @lines.size > @height || (@lines.size == @height && @col > 0)
+          @overflow_callback.call(@lines.take(@height))
+          @multipage = true
+        end
+      end
+
+      def multipage?
+        @multipage
+      end
+
+      alias print write
+      alias << write
+    end
   end
 end

test/irb/test_context.rb

@@ -361,7 +361,7 @@ def test_omit_multiline_on_assignment
           irb.eval_input
         end
         assert_empty err
-        assert_equal("=> #{value_first_line[0..(input.winsize.last - 9)]}...\n=> \n#{value}\n", out)
+        assert_equal("=> \n#{value_first_line[0, input.winsize.last]}...\n=> \n#{value}\n", out)
         irb.context.evaluate_expression('A.remove_method(:inspect)', 0)
 
         input.reset

test/irb/test_pager.rb

@@ -0,0 +1,47 @@
+# frozen_string_literal: false
+require 'irb/pager'
+
+require_relative 'helper'
+
+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' * 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; }
+    end
+  end
+
+  class PageOverflowIOTest < TestCase
+    def test_overflow
+      actual_events = []
+      out = IRB::Pager::PageOverflowIO.new(10, 4) do |lines|
+        actual_events << [:callback_called, lines]
+      end
+      out.puts 'a' * 15
+      out.write  'b' * 15
+
+      actual_events << :before_write
+      out.write 'c' * 1000
+      actual_events << :after_write
+
+      out.puts 'd' * 1000
+      out.write 'e' * 1000
+
+      expected_events = [
+        :before_write,
+        [:callback_called, ['a' * 10, 'a' * 5 + "\n",  'b' * 10, 'b' * 5 + 'c' * 5]],
+        :after_write,
+      ]
+      assert_equal expected_events, actual_events
+
+      expected_whole_content = 'a' * 15 + "\n" + 'b' * 15 + 'c' * 1000 + 'd' * 1000 + "\n" + 'e' * 1000
+      assert_equal expected_whole_content, out.string
+    end
+  end
+end

test/irb/yamatanooroti/test_rendering.rb

@@ -385,12 +385,28 @@ def test_long_evaluation_output_is_paged
     write("'a' * 80 * 11\n")
     write("'foo' + 'bar'\n") # eval something to make sure IRB resumes
 
-    assert_screen(/(a{80}\n){8}/)
+    assert_screen(/"a{79}\n(a{80}\n){7}/)
     # because pager is invoked, foobar will not be evaluated
     assert_screen(/\A(?!foobar)/)
     close
   end
 
+  def test_pretty_print_preview_with_slow_inspect
+    write_irbrc <<~'LINES'
+      require "irb/pager"
+    LINES
+    start_terminal(10, 80, %W{ruby -I#{@pwd}/lib #{@pwd}/exe/irb}, startup_message: /irb\(main\)/)
+    write("o1 = Object.new; def o1.inspect; 'INSPECT'; end\n")
+    write("o2 = Object.new; def o2.inspect; sleep 10; end\n")
+    # preview should be shown even if pretty_print is not completed.
+    write("[o1] * 20 + [o2]\n")
+    assert_screen(/=>\n\[INSPECT,\n( INSPECT,\n){7}/)
+    write("\C-c") # abort pretty_print
+    write("'foo' + 'bar'\n") # eval something to make sure IRB resumes
+    assert_screen(/foobar/)
+    close
+  end
+
   def test_long_evaluation_output_is_preserved_after_paging
     write_irbrc <<~'LINES'
       require "irb/pager"