From 5b61660ad70bd63a21ccb8d31fb21c1d82d6cff6 Mon Sep 17 00:00:00 2001
From: Marc-Andre Lafortune <github@marc-andre.ca>
Date: Thu, 23 Jul 2020 23:48:57 -0400
Subject: [PATCH] Add `Parser::Source::Range#line_span`

---
 CHANGELOG.md                       |  1 +
 lib/rubocop/ast.rb                 |  1 +
 lib/rubocop/ast/ext/range.rb       | 28 ++++++++++++++++++++++++++++
 spec/rubocop/ast/ext/range_spec.rb | 22 ++++++++++++++++++++++
 4 files changed, 52 insertions(+)
 create mode 100644 lib/rubocop/ast/ext/range.rb
 create mode 100644 spec/rubocop/ast/ext/range_spec.rb

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 880c89753..ec5096461 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,7 @@
 * [#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][])
+* [#84](https://github.com/rubocop-hq/rubocop-ast/pull/84): Add `Source::Range#line_span` ([@marcandre][])
 
 ### Bug fixes
 
diff --git a/lib/rubocop/ast.rb b/lib/rubocop/ast.rb
index dd9f6bcff..503e525fa 100644
--- a/lib/rubocop/ast.rb
+++ b/lib/rubocop/ast.rb
@@ -4,6 +4,7 @@
 require 'forwardable'
 require 'set'
 
+require_relative 'ast/ext/range'
 require_relative 'ast/node_pattern'
 require_relative 'ast/sexp'
 require_relative 'ast/node'
diff --git a/lib/rubocop/ast/ext/range.rb b/lib/rubocop/ast/ext/range.rb
new file mode 100644
index 000000000..b2daa0815
--- /dev/null
+++ b/lib/rubocop/ast/ext/range.rb
@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+module RuboCop
+  module AST
+    module Ext
+      # Extensions to Parser::AST::Range
+      module Range
+        # @return [Range] the range of line numbers for the node
+        # If `exclude_end` is `true`, then the range will be exclusive.
+        #
+        # Assume that `node` corresponds to the following array literal:
+        #
+        #   [
+        #     :foo,
+        #     :bar
+        #   ]
+        #
+        #   node.loc.begin.line_span                         # => 1..1
+        #   node.loc.expression.line_span(exclude_end: true) # => 1...4
+        def line_span(exclude_end: false)
+          ::Range.new(first_line, last_line, exclude_end)
+        end
+      end
+    end
+  end
+end
+
+::Parser::Source::Range.include ::RuboCop::AST::Ext::Range
diff --git a/spec/rubocop/ast/ext/range_spec.rb b/spec/rubocop/ast/ext/range_spec.rb
new file mode 100644
index 000000000..86d82bff8
--- /dev/null
+++ b/spec/rubocop/ast/ext/range_spec.rb
@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+RSpec.describe RuboCop::AST::Ext::Range do
+  let(:source) { <<~RUBY }
+    [
+      1,
+      2
+    ]
+  RUBY
+
+  let(:node) { parse_source(source).ast }
+
+  describe '#line_span' do
+    it 'returns the range of lines a range occupies' do
+      expect(node.loc.begin.line_span).to eq 1..1
+    end
+
+    it 'accepts an `exclude_end` keyword argument' do
+      expect(node.loc.expression.line_span(exclude_end: true)).to eq 1...4
+    end
+  end
+end