Add measure command

lib/irb.rb

@@ -538,7 +538,23 @@ def eval_input
         signal_status(:IN_EVAL) do
           begin
             line.untaint if RUBY_VERSION < '2.7'
-            @context.evaluate(line, line_no, exception: exc)
+            if IRB.conf[:MEASURE] && IRB.conf[:MEASURE_CALLBACKS].empty?
+              IRB.set_measure_callback
+            end
+            if IRB.conf[:MEASURE] && !IRB.conf[:MEASURE_CALLBACKS].empty?
+              result = nil
+              last_proc = proc{ result = @context.evaluate(line, line_no, exception: exc) }
+              IRB.conf[:MEASURE_CALLBACKS].map{ |s| s.last }.inject(last_proc) { |chain, item|
+                proc {
+                  item.(@context, line, line_no, exception: exc) do
+                    chain.call
+                  end
+                }
+              }.call
+              @context.set_last_value(result)
+            else
+              @context.evaluate(line, line_no, exception: exc)
+            end
             if @context.echo?
               if assignment_expression?(line)
                 if @context.echo_on_assignment?

lib/irb/cmd/measure.rb

@@ -0,0 +1,34 @@
+require_relative "nop"
+
+# :stopdoc:
+module IRB
+  module ExtendCommand
+    class Measure < Nop
+      def initialize(*args)
+        super(*args)
+      end
+
+      def execute(type = nil, arg = nil)
+        case type
+        when :off
+          IRB.conf[:MEASURE] = nil
+          IRB.unset_measure_callback(arg)
+        when :list
+          IRB.conf[:MEASURE_CALLBACKS].each do |type_name, _|
+            puts "- #{type_name}"
+          end
+        when :on
+          IRB.conf[:MEASURE] = true
+          added = IRB.set_measure_callback(type)
+          puts "#{added.first} is added."
+        else
+          IRB.conf[:MEASURE] = true
+          added = IRB.set_measure_callback(type)
+          puts "#{added.first} is added."
+        end
+        nil
+      end
+    end
+  end
+end
+# :startdoc:

lib/irb/extend-command.rb

@@ -125,6 +125,10 @@ def irb_context
         :irb_info, :Info, "irb/cmd/info"
       ],
 
+      [
+        :measure, :Measure, "irb/cmd/measure"
+      ],
+
     ]
 
     # Installs the default irb commands:

lib/irb/init.rb

@@ -111,11 +111,66 @@ def IRB.init_config(ap_path)
     @CONF[:CONTEXT_MODE] = 4 # use a copy of TOPLEVEL_BINDING
     @CONF[:SINGLE_IRB] = false
 
+    @CONF[:MEASURE] = false
+    @CONF[:MEASURE_PROC] = {}
+    @CONF[:MEASURE_PROC][:TIME] = proc { |context, code, line_no, &block|
+      time = Time.now
+      result = block.()
+      now = Time.now
+      puts 'processing time: %fs' % (now - time) if IRB.conf[:MEASURE]
+      result
+    }
+    @CONF[:MEASURE_PROC][:STACKPROF] = proc { |context, code, line_no, &block|
+      success = false
+      begin
+        require 'stackprof'
+        success = true
+      rescue LoadError
+        puts 'Please run "gem install stackprof" before measuring by StackProf.'
+      end
+      if success
+        result = nil
+        stackprof_result = StackProf.run(mode: :cpu) do
+          result = block.()
+        end
+        StackProf::Report.new(stackprof_result).print_text if IRB.conf[:MEASURE]
+        result
+      else
+        block.()
+      end
+    }
+    @CONF[:MEASURE_CALLBACKS] = []
+
     @CONF[:LC_MESSAGES] = Locale.new
 
     @CONF[:AT_EXIT] = []
   end
 
+  def IRB.set_measure_callback(type = nil)
+    added = nil
+    if type
+      type_sym = type.upcase.to_sym
+      if IRB.conf[:MEASURE_PROC][type_sym]
+        added = [type_sym, IRB.conf[:MEASURE_PROC][type_sym]]
+      end
+    elsif IRB.conf[:MEASURE_PROC][:CUSTOM]
+      added = [:CUSTOM, IRB.conf[:MEASURE_PROC][:CUSTOM]]
+    else
+      added = [:TIME, IRB.conf[:MEASURE_PROC][:TIME]]
+    end
+    IRB.conf[:MEASURE_CALLBACKS] << added if added
+    added
+  end
+
+  def IRB.unset_measure_callback(type = nil)
+    if type.nil?
+      IRB.conf[:MEASURE_CALLBACKS].clear
+    else
+      type_sym = type.upcase.to_sym
+      IRB.conf[:MEASURE_CALLBACKS].reject!{ |t, c| t == type_sym }
+    end
+  end
+
   def IRB.init_error
     @CONF[:LC_MESSAGES].load("irb/error.rb")
   end

test/irb/test_cmd.rb

@@ -123,5 +123,149 @@ def test_irb_info_singleline_without_rc_files
       IRB.__send__(:remove_const, :IRBRC_EXT)
       IRB.const_set(:IRBRC_EXT, ext_backup)
     end
+
+    class TestInputMethod < ::IRB::InputMethod
+      attr_reader :list, :line_no
+
+      def initialize(list = [])
+        super("test")
+        @line_no = 0
+        @list = list
+      end
+
+      def gets
+        @list[@line_no]&.tap {@line_no += 1}
+      end
+
+      def eof?
+        @line_no >= @list.size
+      end
+
+      def encoding
+        Encoding.default_external
+      end
+
+      def reset
+        @line_no = 0
+      end
+    end
+
+    def test_measure
+      IRB.init_config(nil)
+      IRB.conf[:PROMPT] = {
+        DEFAULT: {
+          PROMPT_I: '> ',
+          PROMPT_S: '> ',
+          PROMPT_C: '> ',
+          PROMPT_N: '> '
+        }
+      }
+      IRB.conf[:MEASURE] = false
+      input = TestInputMethod.new([
+        "3\n",
+        "measure\n",
+        "3\n",
+        "measure :off\n",
+        "3\n",
+      ])
+      irb = IRB::Irb.new(IRB::WorkSpace.new(Object.new), input)
+      irb.context.return_format = "=> %s\n"
+      out, err = capture_output do
+        irb.eval_input
+      end
+      assert_empty err
+      assert_match(/\A=> 3\nTIME is added\.\n=> nil\nprocessing time: .+\n=> 3\n=> nil\n=> 3\n/, out)
+    end
+
+    def test_measure_enabled_by_rc
+      IRB.init_config(nil)
+      IRB.conf[:PROMPT] = {
+        DEFAULT: {
+          PROMPT_I: '> ',
+          PROMPT_S: '> ',
+          PROMPT_C: '> ',
+          PROMPT_N: '> '
+        }
+      }
+      IRB.conf[:MEASURE] = true
+      input = TestInputMethod.new([
+        "3\n",
+        "measure :off\n",
+        "3\n",
+      ])
+      irb = IRB::Irb.new(IRB::WorkSpace.new(Object.new), input)
+      irb.context.return_format = "=> %s\n"
+      out, err = capture_output do
+        irb.eval_input
+      end
+      assert_empty err
+      assert_match(/\Aprocessing time: .+\n=> 3\n=> nil\n=> 3\n/, out)
+    end
+
+    def test_measure_enabled_by_rc_with_custom
+      IRB.init_config(nil)
+      IRB.conf[:PROMPT] = {
+        DEFAULT: {
+          PROMPT_I: '> ',
+          PROMPT_S: '> ',
+          PROMPT_C: '> ',
+          PROMPT_N: '> '
+        }
+      }
+      IRB.conf[:MEASURE] = true
+      IRB.conf[:MEASURE_PROC][:CUSTOM] = proc { |line, line_no, &block|
+        time = Time.now
+        result = block.()
+        now = Time.now
+        puts 'custom processing time: %fs' % (Time.now - time) if IRB.conf[:MEASURE]
+        result
+      }
+      input = TestInputMethod.new([
+        "3\n",
+        "measure :off\n",
+        "3\n",
+      ])
+      irb = IRB::Irb.new(IRB::WorkSpace.new(Object.new), input)
+      irb.context.return_format = "=> %s\n"
+      out, err = capture_output do
+        irb.eval_input
+      end
+      assert_empty err
+      assert_match(/\Acustom processing time: .+\n=> 3\n=> nil\n=> 3\n/, out)
+    end
+
+    def test_measure_with_custom
+      IRB.init_config(nil)
+      IRB.conf[:PROMPT] = {
+        DEFAULT: {
+          PROMPT_I: '> ',
+          PROMPT_S: '> ',
+          PROMPT_C: '> ',
+          PROMPT_N: '> '
+        }
+      }
+      IRB.conf[:MEASURE] = false
+      IRB.conf[:MEASURE_PROC][:CUSTOM] = proc { |line, line_no, &block|
+        time = Time.now
+        result = block.()
+        now = Time.now
+        puts 'custom processing time: %fs' % (Time.now - time) if IRB.conf[:MEASURE]
+        result
+      }
+      input = TestInputMethod.new([
+        "3\n",
+        "measure\n",
+        "3\n",
+        "measure :off\n",
+        "3\n",
+      ])
+      irb = IRB::Irb.new(IRB::WorkSpace.new(Object.new), input)
+      irb.context.return_format = "=> %s\n"
+      out, err = capture_output do
+        irb.eval_input
+      end
+      assert_empty err
+      assert_match(/\A=> 3\nCUSTOM is added\.\n=> nil\ncustom processing time: .+\n=> 3\n=> nil\n=> 3\n/, out)
+    end
   end
 end