Merge pull request #291 from kaiquekandykoga/heredoc-indent

Improve HEREDOCS
ruby · Sep 22, 2021 · 360c9c3 · 360c9c3
2 parents 2952916 + 6fec2a5
commit 360c9c3
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 10 deletions.
diff --git a/lib/irb/ruby-lex.rb b/lib/irb/ruby-lex.rb
@@ -207,7 +207,7 @@ def set_auto_indent(context)
           last_line = lines[line_index]&.byteslice(0, byte_pointer)
           code += last_line if last_line
           @tokens = self.class.ripper_lex_without_warning(code, context: context)
-          corresponding_token_depth = check_corresponding_token_depth
+          corresponding_token_depth = check_corresponding_token_depth(lines, line_index)
           if corresponding_token_depth
             corresponding_token_depth
           else
@@ -603,14 +603,19 @@ def check_newline_depth_difference
     depth_difference
   end
 
-  def check_corresponding_token_depth
+  def check_corresponding_token_depth(lines, line_index)
     corresponding_token_depth = nil
     is_first_spaces_of_line = true
     is_first_printable_of_line = true
     spaces_of_nest = []
     spaces_at_line_head = 0
     open_brace_on_line = 0
     in_oneliner_def = nil
+
+    if heredoc_scope?
+      return lines[line_index][/^ */].length
+    end
+
     @tokens.each_with_index do |t, index|
       # detecting one-liner method definition
       if in_oneliner_def.nil?
@@ -817,5 +822,12 @@ def check_termination_in_prev_line(code, context: nil)
     end
     false
   end
+
+  private
+
+  def heredoc_scope?
+    heredoc_tokens = @tokens.select { |t| [:on_heredoc_beg, :on_heredoc_end].include?(t.event) }
+    heredoc_tokens[-1]&.event == :on_heredoc_beg
+  end
 end
 # :startdoc:
diff --git a/test/irb/test_ruby_lex.rb b/test/irb/test_ruby_lex.rb
@@ -109,8 +109,8 @@ def test_multiple_braces_in_a_line
         Row.new(%q(    ]), 4, 4),
         Row.new(%q(  ]), 2, 2),
         Row.new(%q(]), 0, 0),
-        Row.new(%q([<<FOO]), nil, 0),
-        Row.new(%q(hello), nil, 0),
+        Row.new(%q([<<FOO]), 0, 0),
+        Row.new(%q(hello), 0, 0),
         Row.new(%q(FOO), nil, 0),
       ]
 
@@ -465,10 +465,10 @@ def test_do_corresponding_to_loop
 
     def test_heredoc_with_indent
       input_with_correct_indents = [
-        Row.new(%q(<<~Q), nil, 0, 0),
-        Row.new(%q({), nil, 0, 0),
-        Row.new(%q(  #), nil, 0, 0),
-        Row.new(%q(}), nil, 0, 0),
+        Row.new(%q(<<~Q), 0, 0, 0),
+        Row.new(%q({), 0, 0, 0),
+        Row.new(%q(  #), 2, 0, 0),
+        Row.new(%q(}), 0, 0, 0)
       ]
 
       lines = []
@@ -503,8 +503,8 @@ def test_broken_heredoc
       end
       input_with_correct_indents = [
         Row.new(%q(def foo), nil, 2, 1),
-        Row.new(%q(  <<~Q), nil, 2, 1),
-        Row.new(%q(  Qend), nil, 2, 1),
+        Row.new(%q(  <<~Q), 2, 2, 1),
+        Row.new(%q(  Qend), 2, 2, 1),
       ]
 
       lines = []