Allow "measure" command to take block

ruby · Feb 6, 2021 · 20f1ca2 · 20f1ca2
1 parent 7248da8
commit 20f1ca2
Show file tree

Hide file tree

Showing 4 changed files with 49 additions and 7 deletions.
diff --git a/lib/irb/cmd/measure.rb b/lib/irb/cmd/measure.rb
@@ -8,7 +8,7 @@ def initialize(*args)
         super(*args)
       end
 
-      def execute(type = nil, arg = nil)
+      def execute(type = nil, arg = nil, &block)
         case type
         when :off
           IRB.conf[:MEASURE] = nil
@@ -22,9 +22,15 @@ def execute(type = nil, arg = nil)
           added = IRB.set_measure_callback(type, arg)
           puts "#{added[0]} is added." if added
         else
-          IRB.conf[:MEASURE] = true
-          added = IRB.set_measure_callback(type, arg)
-          puts "#{added[0]} is added." if added
+          if block_given?
+            IRB.conf[:MEASURE] = true
+            added = IRB.set_measure_callback(&block)
+            puts "#{added[0]} is added." if added
+          else
+            IRB.conf[:MEASURE] = true
+            added = IRB.set_measure_callback(type, arg)
+            puts "#{added[0]} is added." if added
+          end
         end
         nil
       end

diff --git a/lib/irb/cmd/nop.rb b/lib/irb/cmd/nop.rb
@@ -15,9 +15,9 @@ module ExtendCommand
     class Nop
 
 
-      def self.execute(conf, *opts)
+      def self.execute(conf, *opts, &block)
         command = new(conf)
-        command.execute(*opts)
+        command.execute(*opts, &block)
       end
 
       def initialize(conf)

diff --git a/lib/irb/init.rb b/lib/irb/init.rb
@@ -146,7 +146,7 @@ def IRB.init_config(ap_path)
     @CONF[:AT_EXIT] = []
   end
 
-  def IRB.set_measure_callback(type = nil, arg = nil)
+  def IRB.set_measure_callback(type = nil, arg = nil, &block)
     added = nil
     if type
       type_sym = type.upcase.to_sym
@@ -155,6 +155,8 @@ def IRB.set_measure_callback(type = nil, arg = nil)
       end
     elsif IRB.conf[:MEASURE_PROC][:CUSTOM]
       added = [:CUSTOM, IRB.conf[:MEASURE_PROC][:CUSTOM], arg]
+    elsif block_given?
+      added = [:BLOCK, block, arg]
     else
       added = [:TIME, IRB.conf[:MEASURE_PROC][:TIME], arg]
     end

diff --git a/test/irb/test_cmd.rb b/test/irb/test_cmd.rb
@@ -276,6 +276,40 @@ def test_measure_with_custom
       assert_match(/\A=> 3\nCUSTOM is added\.\n=> nil\ncustom processing time: .+\n=> 3\n=> nil\n=> 3\n/, out)
     end
 
+    def test_measure_with_proc
+      IRB.init_config(nil)
+      IRB.conf[:PROMPT] = {
+        DEFAULT: {
+          PROMPT_I: '> ',
+          PROMPT_S: '> ',
+          PROMPT_C: '> ',
+          PROMPT_N: '> '
+        }
+      }
+      IRB.conf[:VERBOSE] = false
+      IRB.conf[:PROMPT_MODE] = :DEFAULT
+      IRB.conf[:MEASURE] = false
+      input = TestInputMethod.new([
+        "3\n",
+        "measure { |context, code, line_no, &block|\n",
+        "  result = block.()\n",
+        "  puts 'aaa' if IRB.conf[:MEASURE]\n",
+        "}\n",
+        "3\n",
+        "measure :off\n",
+        "3\n",
+      ])
+      c = Class.new(Object)
+      irb = IRB::Irb.new(IRB::WorkSpace.new(c.new), input)
+      irb.context.return_format = "=> %s\n"
+      out, err = capture_output do
+        irb.eval_input
+      end
+      assert_empty err
+      assert_match(/\A=> 3\nBLOCK is added\.\n=> nil\naaa\n=> 3\n=> nil\n=> 3\n/, out)
+      assert_empty(c.class_variables)
+    end
+
     def test_irb_source
       IRB.init_config(nil)
       File.write("#{@tmpdir}/a.rb", "a = 'hi'\n")