Add BeginNode for kwbegin nodes.

changelog/change_add_beginnode_for_kwbegin_nodes.md

@@ -0,0 +1 @@
+* [#333](https://github.com/rubocop/rubocop-ast/pull/333): Add `EnsureNode#rescue_node` method. ([@dvandersluis][])

changelog/new_add_beginnode_for_kwbegin_nodes.md

@@ -0,0 +1 @@
+* [#333](https://github.com/rubocop/rubocop-ast/pull/333): Add `BeginNode` for `kwbegin` nodes. ([@dvandersluis][])

docs/modules/ROOT/pages/node_types.adoc

@@ -142,7 +142,7 @@ The following fields are given when relevant to nodes in the source code:
 
 |kwarg|Required keyword argument. Must come inside an `args`.|One child - a symbol, representing the argument name.|def foo(bar:)|https://rubydoc.info/github/rubocop/rubocop-ast/RuboCop/AST/ArgNode[ArgNode]
 
-|kwbegin|Explicit `begin` block.|Child nodes are body statements.|begin,end|N/A
+|kwbegin|Explicit `begin` block.|Child nodes are body statements.|begin,end|https://rubydoc.info/github/rubocop/rubocop-ast/RuboCop/AST/KeywordBeginNode[KeywordBeginNode]
 
 |kwnilarg|Double splat with nil in function definition, used to specify that the function does not accept keyword args. Must come inside an `args`.|None|def foo(**nil)|N/A
 

lib/rubocop/ast.rb

@@ -60,6 +60,7 @@
 require_relative 'ast/node/index_node'
 require_relative 'ast/node/indexasgn_node'
 require_relative 'ast/node/int_node'
+require_relative 'ast/node/keyword_begin_node'
 require_relative 'ast/node/keyword_splat_node'
 require_relative 'ast/node/lambda_node'
 require_relative 'ast/node/masgn_node'

lib/rubocop/ast/builder.rb

@@ -63,6 +63,7 @@ class Builder < Parser::Builders::Default
         irange:              RangeNode,
         erange:              RangeNode,
         kwargs:              HashNode,
+        kwbegin:             KeywordBeginNode,
         kwsplat:             KeywordSplatNode,
         lambda:              LambdaNode,
         masgn:               MasgnNode,

lib/rubocop/ast/node/keyword_begin_node.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+module RuboCop
+  module AST
+    # A node extension for `kwbegin` nodes. This will be used in place of a plain
+    # node when the builder constructs the AST, making its methods available
+    # to all `kwbegin` nodes within RuboCop.
+    class KeywordBeginNode < Node
+      # Returns the body of the `kwbegin` block. Returns `self` if the `kwbegin` contains
+      # multiple nodes.
+      #
+      # @return [Node, nil] The body of the `kwbegin`.
+      def body
+        return unless node_parts.any?
+
+        if rescue_node
+          rescue_node.body
+        elsif ensure_node
+          ensure_node.node_parts[0]
+        elsif node_parts.one?
+          node_parts[0]
+        else
+          self
+        end
+      end
+
+      # Returns the `rescue` node of the `kwbegin` block, if one is present.
+      #
+      # @return [Node, nil] The `rescue` node within `kwbegin`.
+      def ensure_node
+        node_parts[0] if node_parts[0]&.ensure_type?
+      end
+
+      # Returns the `rescue` node of the `kwbegin` block, if one is present.
+      #
+      # @return [Node, nil] The `rescue` node within `kwbegin`.
+      def rescue_node
+        return ensure_node&.rescue_node if ensure_node&.rescue_node
+
+        node_parts[0] if node_parts[0]&.rescue_type?
+      end
+    end
+  end
+end

spec/rubocop/ast/keyword_begin_node_spec.rb

@@ -0,0 +1,274 @@
+# frozen_string_literal: true
+
+RSpec.describe RuboCop::AST::KeywordBeginNode do
+  let(:parsed_source) { parse_source(source) }
+  let(:kwbegin_node) { parsed_source.ast }
+
+  describe '.new' do
+    let(:source) do
+      <<~RUBY
+        begin
+          foo
+        end
+      RUBY
+    end
+
+    it { expect(kwbegin_node).to be_a(described_class) }
+  end
+
+  describe '#body' do
+    subject(:body) { kwbegin_node.body }
+
+    let(:node) { parsed_source.node }
+
+    context 'when the `kwbegin` node is empty' do
+      let(:source) do
+        <<~RUBY
+          begin
+          end
+        RUBY
+      end
+
+      it { is_expected.to be_nil }
+    end
+
+    context 'when the `kwbegin` node only contains a single line' do
+      let(:source) do
+        <<~RUBY
+          begin
+            >> foo <<
+          end
+        RUBY
+      end
+
+      it { is_expected.to eq(node) }
+    end
+
+    context 'when the body has multiple lines' do
+      let(:source) do
+        <<~RUBY
+          begin
+            foo
+            bar
+          end
+        RUBY
+      end
+
+      it 'returns the entire `kwbegin` node' do
+        expect(body).to eq(kwbegin_node)
+      end
+    end
+
+    context 'when there is a `rescue` node' do
+      let(:source) do
+        <<~RUBY
+          begin
+            >>foo<<
+          rescue
+            bar
+          end
+        RUBY
+      end
+
+      it { is_expected.to eq(node) }
+    end
+
+    context 'when there is an `ensure` node' do
+      let(:source) do
+        <<~RUBY
+          begin
+            >>foo<<
+          ensure
+            bar
+          end
+        RUBY
+      end
+
+      it { is_expected.to eq(node) }
+    end
+
+    context 'when there is a `rescue` and `ensure` node' do
+      let(:source) do
+        <<~RUBY
+          begin
+            >>foo<<
+          rescue
+            bar
+          ensure
+            baz
+          end
+        RUBY
+      end
+
+      it { is_expected.to eq(node) }
+    end
+  end
+
+  describe '#ensure_node' do
+    subject(:ensure_node) { kwbegin_node.ensure_node }
+
+    context 'when the `kwbegin` node is empty' do
+      let(:source) do
+        <<~RUBY
+          begin
+          end
+        RUBY
+      end
+
+      it { is_expected.to be_nil }
+    end
+
+    context 'when the `kwbegin` node only contains a single line' do
+      let(:source) do
+        <<~RUBY
+          begin
+            foo
+          end
+        RUBY
+      end
+
+      it { is_expected.to be_nil }
+    end
+
+    context 'when the body has multiple lines' do
+      let(:source) do
+        <<~RUBY
+          begin
+            foo
+            bar
+          end
+        RUBY
+      end
+
+      it { is_expected.to be_nil }
+    end
+
+    context 'when there is a `rescue` node without `ensure`' do
+      let(:source) do
+        <<~RUBY
+          begin
+            foo
+          rescue
+            bar
+          end
+        RUBY
+      end
+
+      it { is_expected.to be_nil }
+    end
+
+    context 'when there is an `ensure` node' do
+      let(:source) do
+        <<~RUBY
+          begin
+            foo
+          ensure
+            bar
+          end
+        RUBY
+      end
+
+      it { is_expected.to be_a(RuboCop::AST::EnsureNode) }
+    end
+
+    context 'when there is a `rescue` and `ensure` node' do
+      let(:source) do
+        <<~RUBY
+          begin
+            foo
+          rescue
+            bar
+          ensure
+            baz
+          end
+        RUBY
+      end
+
+      it { is_expected.to be_a(RuboCop::AST::EnsureNode) }
+    end
+  end
+
+  describe '#rescue_node' do
+    subject(:rescue_node) { kwbegin_node.rescue_node }
+
+    context 'when the `kwbegin` node is empty' do
+      let(:source) do
+        <<~RUBY
+          begin
+          end
+        RUBY
+      end
+
+      it { is_expected.to be_nil }
+    end
+
+    context 'when the `kwbegin` node only contains a single line' do
+      let(:source) do
+        <<~RUBY
+          begin
+            foo
+          end
+        RUBY
+      end
+
+      it { is_expected.to be_nil }
+    end
+
+    context 'when the body has multiple lines' do
+      let(:source) do
+        <<~RUBY
+          begin
+            foo
+            bar
+          end
+        RUBY
+      end
+
+      it { is_expected.to be_nil }
+    end
+
+    context 'when there is a `rescue` node without `ensure`' do
+      let(:source) do
+        <<~RUBY
+          begin
+            foo
+          rescue
+            bar
+          end
+        RUBY
+      end
+
+      it { is_expected.to be_a(RuboCop::AST::RescueNode) }
+    end
+
+    context 'when there is an `ensure` node without `rescue`' do
+      let(:source) do
+        <<~RUBY
+          begin
+            foo
+          ensure
+            bar
+          end
+        RUBY
+      end
+
+      it { is_expected.to be_nil }
+    end
+
+    context 'when there is a `rescue` and `ensure` node' do
+      let(:source) do
+        <<~RUBY
+          begin
+            foo
+          rescue
+            bar
+          ensure
+            baz
+          end
+        RUBY
+      end
+
+      it { is_expected.to be_a(RuboCop::AST::RescueNode) }
+    end
+  end
+end