Review comments API

CHANGELOG.md

@@ -7,6 +7,8 @@
 * [#70](https://github.com/rubocop-hq/rubocop-ast/pull/70): Add `NextNode` ([@marcandre][])
 * [#85](https://github.com/rubocop-hq/rubocop-ast/pull/85): Add `IntNode#value` and `FloatNode#value`. ([@fatkodima][])
 * [#82](https://github.com/rubocop-hq/rubocop-ast/pull/82): `NodePattern`: Allow comments ([@marcandre][])
+* [#83](https://github.com/rubocop-hq/rubocop-ast/pull/83): Add `ProcessedSource#comment_at_line` ([@marcandre][])
+* [#83](https://github.com/rubocop-hq/rubocop-ast/pull/83): Add `ProcessedSource#each_comment_in_lines` ([@marcandre][])
 
 ### Bug fixes
 

lib/rubocop/ast/processed_source.rb

@@ -76,7 +76,7 @@ def each_comment(&block)
         comments.each(&block)
       end
 
-      # @deprecated Use `comments.find`
+      # @deprecated Use `comment_at_line`, `each_comment_in_lines`, or `comments.find`
       def find_comment(&block)
         comments.find(&block)
       end
@@ -99,22 +99,39 @@ def blank?
         ast.nil?
       end
 
+      # @return [Comment, nil] the comment at that line, if any.
+      def comment_at_line(line)
+        comment_index[line]
+      end
+
       # @return [Boolean] if the given line number has a comment.
       def line_with_comment?(line)
-        comment_lines.include?(line)
+        comment_index.include?(line)
+      end
+
+      # Enumerates on the comments contained with the given `line_range`
+      def each_comment_in_lines(line_range)
+        return to_enum(:each_comment_in_lines, line_range) unless block_given?
+
+        line_range.each do |line|
+          if (comment = comment_index[line])
+            yield comment
+          end
+        end
       end
 
       # @return [Boolean] if any of the lines in the given `source_range` has a comment.
+      # Consider using `each_comment_in_lines` instead
       def contains_comment?(source_range)
-        (source_range.line..source_range.last_line).any? do |line|
-          line_with_comment?(line)
-        end
+        each_comment_in_lines(source_range.line..source_range.last_line).any?
       end
       # @deprecated use contains_comment?
       alias commented? contains_comment?
 
+      # @deprecated Use `each_comment_in_lines`
+      # Should have been called `comments_before_or_at_line`. Doubtful it has of any valid use.
       def comments_before_line(line)
-        comments.select { |c| c.location.line <= line }
+        each_comment_in_lines(0..line).to_a
       end
 
       def start_with?(string)
@@ -144,8 +161,10 @@ def line_indentation(line_number)
 
       private
 
-      def comment_lines
-        @comment_lines ||= comments.map { |c| c.location.line }
+      def comment_index
+        @comment_index ||= {}.tap do |hash|
+          comments.each { |c| hash[c.location.line] = c }
+        end
       end
 
       def parse(source, ruby_version)

spec/rubocop/ast/processed_source_spec.rb

@@ -218,10 +218,12 @@ def some_method
 
   context 'with heavily commented source' do
     let(:source) { <<~RUBY }
-      def foo # comment one
-        bar # comment two
-      end # comment three
-      foo
+      # comment one
+      [ 1,
+        { a: 2,
+          b: 3 # comment two
+        }
+      ]
     RUBY
 
     describe '#each_comment' do
@@ -233,17 +235,17 @@ def foo # comment one
           comments << item
         end
 
-        expect(comments.size).to eq 3
+        expect(comments.size).to eq 2
       end
     end
 
     describe '#find_comment' do
       it 'yields correct comment' do
         comment = processed_source.find_comment do |item|
-          item.text == '# comment three'
+          item.text == '# comment two'
         end
 
-        expect(comment.text).to eq '# comment three'
+        expect(comment.text).to eq '# comment two'
       end
 
       it 'yields nil when there is no match' do
@@ -255,37 +257,41 @@ def foo # comment one
       end
     end
 
-    describe '#line_with_comment?' do
-      let(:source) { <<~RUBY }
-        # comment
-        [
-          1, # comment
-          2
-        ]
-      RUBY
+    describe '#comment_at_line' do
+      it 'returns the comment at the given line number' do
+        expect(processed_source.comment_at_line(1).text).to eq '# comment one'
+        expect(processed_source.comment_at_line(4).text).to eq '# comment two'
+      end
+
+      it 'returns nil if line has no comment' do
+        expect(processed_source.comment_at_line(3)).to be nil
+      end
+    end
 
+    describe '#each_comment_in_lines' do
+      it 'yields the comments' do
+        enum = processed_source.each_comment_in_lines(1..4)
+        expect(enum.is_a?(Enumerable)).to be(true)
+        expect(enum.to_a).to eq processed_source.comments
+        expect(processed_source.each_comment_in_lines(2..5).map(&:text)).to eq ['# comment two']
+      end
+    end
+
+    describe '#line_with_comment?' do
       it 'returns true for lines with comments' do
         expect(processed_source.line_with_comment?(1)).to be true
-        expect(processed_source.line_with_comment?(3)).to be true
+        expect(processed_source.line_with_comment?(4)).to be true
       end
 
       it 'returns false for lines without comments' do
         expect(processed_source.line_with_comment?(2)).to be false
-        expect(processed_source.line_with_comment?(4)).to be false
+        expect(processed_source.line_with_comment?(5)).to be false
       end
     end
 
     describe '#contains_comment?' do
       subject(:commented) { processed_source.contains_comment?(range) }
 
-      let(:source) { <<~RUBY }
-        # comment
-        [ 1,
-          { a: 2,
-            b: 3 # comment
-          }
-        ]
-      RUBY
       let(:ast) { processed_source.ast }
       let(:array) { ast }
       let(:hash) { array.children[1] }