Add CaseNode#branches

CHANGELOG.md

@@ -10,6 +10,7 @@
 * [#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][])
+* [#87](https://github.com/rubocop-hq/rubocop-ast/pull/87): Add `CaseNode#branches` ([@marcandre][])
 
 ### Bug fixes
 

lib/rubocop/ast/node/case_node.rb

@@ -31,6 +31,16 @@ def when_branches
         node_parts[1...-1]
       end
 
+      # Returns an array of all the when branches in the `case` statement.
+      #
+      # @return [Array<Node, nil>] an array of the bodies of the when branches
+      # and the else (if any). Note that these bodies could be nil.
+      def branches
+        bodies = when_branches.map(&:body)
+        bodies.push(else_branch) if else?
+        bodies
+      end
+
       # Returns the else branch of the `case` statement, if any.
       #
       # @return [Node] the else branch node of the `case` statement

spec/rubocop/ast/case_node_spec.rb

@@ -1,7 +1,8 @@
 # frozen_string_literal: true
 
 RSpec.describe RuboCop::AST::CaseNode do
-  let(:case_node) { parse_source(source).ast }
+  let(:ast) { parse_source(source).ast }
+  let(:case_node) { ast }
 
   describe '.new' do
     let(:source) do
@@ -104,4 +105,47 @@
       end
     end
   end
+
+  describe '#branches' do
+    context 'when there is an else' do
+      let(:source) { <<~RUBY }
+        case
+        when :foo then # do nothing
+        when :bar then 42
+        else 'hello'
+        end
+      RUBY
+
+      it 'returns all the bodies' do
+        expect(case_node.branches).to match [nil, be_int_type, be_str_type]
+      end
+
+      context 'with an empty else' do
+        let(:source) { <<~RUBY }
+          case
+          when :foo then # do nothing
+          when :bar then 42
+          else # do nothing
+          end
+        RUBY
+
+        it 'returns all the bodies' do
+          expect(case_node.branches).to match [nil, be_int_type, nil]
+        end
+      end
+    end
+
+    context 'when there is no else keyword' do
+      let(:source) { <<~RUBY }
+        case
+        when :foo then # do nothing
+        when :bar then 42
+        end
+      RUBY
+
+      it 'returns only then when bodies' do
+        expect(case_node.branches).to match [nil, be_int_type]
+      end
+    end
+  end
 end