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] 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