Warn and do nothing if block is passed to measure command

lib/irb/cmd/measure.rb

@@ -12,32 +12,29 @@ def initialize(*args)
         super(*args)
       end
 
-      def execute(type = nil, arg = nil, &block)
+      def execute(type = nil, arg = nil)
         # Please check IRB.init_config in lib/irb/init.rb that sets
         # IRB.conf[:MEASURE_PROC] to register default "measure" methods,
         # "measure :time" (abbreviated as "measure") and "measure :stackprof".
+
+        if block_given?
+          warn 'Configure IRB.conf[:MEASURE_PROC] to add custom measure methods.'
+          return
+        end
+
         case type
         when :off
-          IRB.conf[:MEASURE] = nil
           IRB.unset_measure_callback(arg)
         when :list
           IRB.conf[:MEASURE_CALLBACKS].each do |type_name, _, arg_val|
             puts "- #{type_name}" + (arg_val ? "(#{arg_val.inspect})" : '')
           end
         when :on
-          IRB.conf[:MEASURE] = true
-          added = IRB.set_measure_callback(type, arg)
+          added = IRB.set_measure_callback(arg)
           puts "#{added[0]} is added." if added
         else
-          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
+          added = IRB.set_measure_callback(type, arg)
+          puts "#{added[0]} is added." if added
         end
         nil
       end

lib/irb/init.rb

@@ -215,6 +215,7 @@ def IRB.set_measure_callback(type = nil, arg = nil, &block)
       added = [:TIME, IRB.conf[:MEASURE_PROC][:TIME], arg]
     end
     if added
+      IRB.conf[:MEASURE] = true
       found = IRB.conf[:MEASURE_CALLBACKS].find{ |m| m[0] == added[0] && m[2] == added[2] }
       if found
         # already added
@@ -235,6 +236,7 @@ def IRB.unset_measure_callback(type = nil)
       type_sym = type.upcase.to_sym
       IRB.conf[:MEASURE_CALLBACKS].reject!{ |t, | t == type_sym }
     end
+    IRB.conf[:MEASURE] = nil if IRB.conf[:MEASURE_CALLBACKS].empty?
   end
 
   def IRB.init_error

test/irb/test_cmd.rb

@@ -226,17 +226,20 @@ def test_measure
 
       c = Class.new(Object)
       out, err = execute_lines(
-        "3\n",
         "measure\n",
         "3\n",
         "measure :off\n",
         "3\n",
+        "measure :on\n",
+        "3\n",
+        "measure :off\n",
+        "3\n",
         conf: conf,
         main: c
       )
 
       assert_empty err
-      assert_match(/\A=> 3\nTIME is added\.\n=> nil\nprocessing time: .+\n=> 3\n=> nil\n=> 3\n/, out)
+      assert_match(/\A(TIME is added\.\n=> nil\nprocessing time: .+\n=> 3\n=> nil\n=> 3\n){2}/, out)
       assert_empty(c.class_variables)
     end
 
@@ -353,7 +356,13 @@ 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
+    def test_measure_toggle
+      measuring_proc = proc { |line, line_no, &block|
+        time = Time.now
+        result = block.()
+        puts 'custom processing time: %fs' % (Time.now - time) if IRB.conf[:MEASURE]
+        result
+      }
       conf = {
         PROMPT: {
           DEFAULT: {
@@ -364,30 +373,49 @@ def test_measure_with_proc
         },
         PROMPT_MODE: :DEFAULT,
         MEASURE: false,
+        MEASURE_PROC: {
+          FOO: proc { |&block| puts 'foo'; block.call },
+          BAR: proc { |&block| puts 'bar'; block.call }
+        }
       }
-      c = Class.new(Object)
       out, err = execute_lines(
+        "measure :foo",
+        "measure :on, :bar",
         "3\n",
-        "measure { |context, code, line_no, &block|\n",
-        "  result = block.()\n",
-        "  puts 'aaa' if IRB.conf[:MEASURE]\n",
-        "  result\n",
-        "}\n",
+        "measure :off, :foo\n",
+        "measure :off, :bar\n",
         "3\n",
-        "measure { |context, code, line_no, &block|\n",
-        "  result = block.()\n",
-        "  puts 'bbb' if IRB.conf[:MEASURE]\n",
-        "  result\n",
-        "}\n",
+        conf: conf
+      )
+
+      assert_empty err
+      assert_match(/\AFOO is added\.\n=> nil\nfoo\nBAR is added\.\n=> nil\nbar\nfoo\n=> 3\nbar\nfoo\n=> nil\nbar\n=> nil\n=> 3\n/, out)
+    end
+
+    def test_measure_with_proc_warning
+      conf = {
+        PROMPT: {
+          DEFAULT: {
+            PROMPT_I: '> ',
+            PROMPT_S: '> ',
+            PROMPT_C: '> '
+          }
+        },
+        PROMPT_MODE: :DEFAULT,
+        MEASURE: false,
+      }
+      c = Class.new(Object)
+      out, err = execute_lines(
         "3\n",
-        "measure :off\n",
+        "measure do\n",
+        "end\n",
         "3\n",
         conf: conf,
         main: c
       )
 
-      assert_empty err
-      assert_match(/\A=> 3\nBLOCK is added\.\n=> nil\naaa\n=> 3\nBLOCK is added.\naaa\n=> nil\nbbb\n=> 3\n=> nil\n=> 3\n/, out)
+      assert_match(/to add custom measure/, err)
+      assert_match(/\A=> 3\n=> nil\n=> 3\n/, out)
       assert_empty(c.class_variables)
     end
   end