From 5e917b18aa05940487878807766cdc416bf67ca6 Mon Sep 17 00:00:00 2001
From: Koichi ITO <koic.ito@gmail.com>
Date: Sun, 11 Jul 2021 16:41:34 +0900
Subject: [PATCH] Add `branches` method for `AST::CaseMatchNode`

This PR adds `branches` method for `AST::CaseMatchNode`.

This makes a method similar to `AST::CaseNode#branches` available
in `AST::CaseMatchNode`.
After this, I will be developing with this API on RuboCop.
---
 ...branches_method_for_ast_case_match_node.md |  1 +
 lib/rubocop/ast/node/case_match_node.rb       | 13 ++++++
 spec/rubocop/ast/case_match_node_spec.rb      | 45 +++++++++++++++++++
 3 files changed, 59 insertions(+)
 create mode 100644 changelog/add_branches_method_for_ast_case_match_node.md

diff --git a/changelog/add_branches_method_for_ast_case_match_node.md b/changelog/add_branches_method_for_ast_case_match_node.md
new file mode 100644
index 000000000..159e31973
--- /dev/null
+++ b/changelog/add_branches_method_for_ast_case_match_node.md
@@ -0,0 +1 @@
+* [#192](https://github.com/rubocop/rubocop-ast/pull/192): Add `branches` method for `AST::CaseMatchNode`. ([@koic][])
diff --git a/lib/rubocop/ast/node/case_match_node.rb b/lib/rubocop/ast/node/case_match_node.rb
index 2b0c6e934..393a5b8d2 100644
--- a/lib/rubocop/ast/node/case_match_node.rb
+++ b/lib/rubocop/ast/node/case_match_node.rb
@@ -31,6 +31,19 @@ def in_pattern_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 `in` branches
+      # and the `else` (if any). Note that these bodies could be nil.
+      def branches
+        bodies = in_pattern_branches.map(&:body)
+        if else?
+          # `empty-else` node sets nil because it has no body.
+          else_branch.empty_else_type? ? bodies.push(nil) : bodies.push(else_branch)
+        end
+        bodies
+      end
+
       # Returns the else branch of the `case` statement, if any.
       #
       # @return [Node] the else branch node of the `case` statement
diff --git a/spec/rubocop/ast/case_match_node_spec.rb b/spec/rubocop/ast/case_match_node_spec.rb
index b5a5286bc..81af40541 100644
--- a/spec/rubocop/ast/case_match_node_spec.rb
+++ b/spec/rubocop/ast/case_match_node_spec.rb
@@ -127,5 +127,50 @@
         end
       end
     end
+
+    describe '#branches' do
+      context 'when there is an else' do
+        context 'with else body' do
+          let(:source) { <<~RUBY }
+            case pattern
+            in :foo then # do nothing
+            in :bar then 42
+            else 'hello'
+            end
+          RUBY
+
+          it 'returns all the bodies' do
+            expect(case_match_node.branches).to match [nil, be_int_type, be_str_type]
+          end
+        end
+
+        context 'with empty else' do
+          let(:source) { <<~RUBY }
+            case pattern
+            in :foo then # do nothing
+            in :bar then 42
+            else # do nothing
+            end
+          RUBY
+
+          it 'returns all the bodies' do
+            expect(case_match_node.branches).to match [nil, be_int_type, nil]
+          end
+        end
+      end
+
+      context 'when there is no else keyword' do
+        let(:source) { <<~RUBY }
+          case pattern
+          in :foo then # do nothing
+          in :bar then 42
+          end
+        RUBY
+
+        it 'returns only then when bodies' do
+          expect(case_match_node.branches).to match [nil, be_int_type]
+        end
+      end
+    end
   end
 end