Use Reline.prompt_proc in IRB

3 files changed, 29 insertions, 6 deletions

diff --git a/lib/irb.rb b/lib/irb.rb
index 1badb9b394..1bab317a05 100644
--- a/lib/irb.rb
+++ b/lib/irb.rb
@@ -469,6 +469,7 @@ module IRB
             @context.io.prompt = p + " " * ind if ind > 0
           end
         end
+        @context.io.prompt
       end
 
       @scanner.set_input(@context.io) do
diff --git a/lib/irb/input-method.rb b/lib/irb/input-method.rb
index aa62a628c8..68d2ad280c 100644
--- a/lib/irb/input-method.rb
+++ b/lib/irb/input-method.rb
@@ -240,12 +240,17 @@ module IRB
       @check_termination_proc = block
     end
 
+    def dynamic_prompt(&block)
+      @prompt_proc = block
+    end
+
     # Reads the next line from this input method.
     #
     # See IO#gets for more information.
     def gets
       Reline.input = @stdin
       Reline.output = @stdout
+      Reline.prompt_proc = @prompt_proc
       if l = readmultiline(@prompt, false, &@check_termination_proc)
         HISTORY.push(l) if !l.empty?
         @line[@line_no += 1] = l + "\n"
diff --git a/lib/irb/ruby-lex.rb b/lib/irb/ruby-lex.rb
index a57a5dfcd9..2c2dfd850a 100644
--- a/lib/irb/ruby-lex.rb
+++ b/lib/irb/ruby-lex.rb
@@ -33,18 +33,26 @@ class RubyLex
     if @io.respond_to?(:check_termination)
       @io.check_termination do |code|
         code.gsub!(/\s*\z/, '').concat("\n")
-        @tokens = Ripper.lex(code)
-        continue = process_continue
-        code_block_open = check_code_block(code)
-        indent = process_nesting_level
-        ltype = process_literal_type
-        if code_block_open or ltype or continue or indent > 0
+        ltype, indent, continue, code_block_open = check_state(code)
+        if ltype or indent > 0 or continue or code_block_open
           false
         else
           true
         end
       end
     end
+    if @io.respond_to?(:dynamic_prompt)
+      @io.dynamic_prompt do |lines, base_line_no|
+        lines << '' if lines.empty?
+        result = []
+        lines.each_index { |i|
+          c = lines[0..i].map{ |l| l + "\n" }.join
+          ltype, indent, continue, code_block_open = check_state(c)
+          result << @prompt.call(ltype, indent, continue, base_line_no + i)
+        }
+        result
+      end
+    end
     if p.respond_to?(:call)
       @input = p
     elsif block_given?
@@ -63,6 +71,15 @@ class RubyLex
     end
   end
 
+  def check_state(code)
+    @tokens = Ripper.lex(code)
+    ltype = process_literal_type
+    indent = process_nesting_level
+    continue = process_continue
+    code_block_open = check_code_block(code)
+    [ltype, indent, continue, code_block_open]
+  end
+
   def prompt
     if @prompt
       @prompt.call(@ltype, @indent, @continue, @line_no)