diff --git a/lib/irb/cmd/measure.rb b/lib/irb/cmd/measure.rb index 5e0bef62af..58eaec2ded 100644 --- a/lib/irb/cmd/measure.rb +++ b/lib/irb/cmd/measure.rb @@ -8,7 +8,7 @@ module IRB 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 @@ module IRB 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 index 9cf4337c28..fa3c011b5f 100644 --- a/lib/irb/cmd/nop.rb +++ b/lib/irb/cmd/nop.rb @@ -15,9 +15,9 @@ module IRB 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 index 8428a4278f..cd57012c61 100644 --- a/lib/irb/init.rb +++ b/lib/irb/init.rb @@ -146,7 +146,7 @@ module IRB # :nodoc: @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 @@ module IRB # :nodoc: 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 index 106e4d2add..045ce0f0e7 100644 --- a/test/irb/test_cmd.rb +++ b/test/irb/test_cmd.rb @@ -276,6 +276,40 @@ module TestIRB 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")