Fix source and lines handling for heredocs

Author: fatkodima

CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### New features
 
+* [#40](https://github.com/rubocop-hq/rubocop-ast/pull/40): Fix source and lines handling for heredocs. ([@fatkodima][])
 * [#37](https://github.com/rubocop-hq/rubocop-ast/pull/37): Add `enumerable_method?` for `MethodIdentifierPredicates`. ([@fatkodima][])
 * [#4](https://github.com/rubocop-hq/rubocop-ast/issues/4): Add `interpolation?` for `RegexpNode`. ([@tejasbubane][])
 * [#20](https://github.com/rubocop-hq/rubocop-ast/pull/20): Add option predicates for `RegexpNode`. ([@owst][])

lib/rubocop/ast/node/str_node.rb

@@ -11,6 +11,40 @@ class StrNode < Node
       def heredoc?
         loc.is_a?(Parser::Source::Map::Heredoc)
       end
+
+      def source
+        if heredoc?
+          source_range.source
+        else
+          super
+        end
+      end
+
+      def source_range
+        if heredoc?
+          Parser::Source::Range.new(loc.expression.source_buffer,
+                                    loc.expression.begin_pos,
+                                    loc.heredoc_end.end_pos)
+        else
+          super
+        end
+      end
+
+      def first_line
+        if heredoc?
+          loc.expression.line
+        else
+          super
+        end
+      end
+
+      def last_line
+        if heredoc?
+          loc.heredoc_end.line
+        else
+          super
+        end
+      end
     end
   end
 end

spec/rubocop/ast/str_node_spec.rb

@@ -56,4 +56,94 @@
       it { expect(str_node.heredoc?).to be(true) }
     end
   end
+
+  describe '#source' do
+    context 'with a normal string' do
+      let(:source) { "'foo'" }
+
+      it { expect(str_node.source).to eq(source) }
+    end
+
+    context 'with a heredoc' do
+      let(:source) do
+        <<~RUBY
+          <<-CODE
+            foo
+            bar
+          CODE
+        RUBY
+      end
+
+      it { expect(str_node.source).to eq(source.rstrip) }
+    end
+  end
+
+  describe '#first_line' do
+    context 'with a normal string' do
+      let(:source) { "'foo'" }
+
+      it { expect(str_node.first_line).to eq(1) }
+    end
+
+    context 'with a heredoc' do
+      let(:source) do
+        <<~RUBY
+          <<-CODE
+            foo
+            bar
+          CODE
+        RUBY
+      end
+
+      it { expect(str_node.first_line).to eq(1) }
+    end
+  end
+
+  describe '#last_line' do
+    context 'with a normal string' do
+      let(:source) do
+        ['"foo"\\',
+         '"bar"'].join("\n")
+      end
+
+      it { expect(str_node.last_line).to eq(2) }
+    end
+
+    context 'with a heredoc' do
+      let(:source) do
+        <<~RUBY
+          <<-CODE
+            foo
+            bar
+          CODE
+        RUBY
+      end
+
+      it { expect(str_node.last_line).to eq(4) }
+    end
+  end
+
+  describe '#line_count' do
+    context 'with a normal string' do
+      let(:source) do
+        ['"foo"\\',
+         '"bar"'].join("\n")
+      end
+
+      it { expect(str_node.line_count).to eq(2) }
+    end
+
+    context 'with a heredoc' do
+      let(:source) do
+        <<~RUBY
+          <<-CODE
+            foo
+            bar
+          CODE
+        RUBY
+      end
+
+      it { expect(str_node.line_count).to eq(4) }
+    end
+  end
 end