diff --git a/lib/rubocop/ast/node/case_node.rb b/lib/rubocop/ast/node/case_node.rb index aa1eae889..ece0cf8d2 100644 --- a/lib/rubocop/ast/node/case_node.rb +++ b/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] 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_branch + 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_node_spec.rb b/spec/rubocop/ast/case_node_spec.rb index 20eb587ad..acb9a60b9 100644 --- a/spec/rubocop/ast/case_node_spec.rb +++ b/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,69 @@ 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] + 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 + + context 'when compared to an IfNode' do + let(:source) { <<~RUBY } + case + when foo then 1 + when bar then 2 + else + end + + if foo then 1 + elsif bar then 2 + else + end + RUBY + + let(:case_node) { ast.children.first } + let(:if_node) { ast.children.last } + + it 'returns the same' do + expect(case_node.branches).to eq if_node.branches + end + end + end end