Add arity to Params, DefNode and BlockNode

CHANGELOG.md

@@ -6,6 +6,10 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/) a
 
 ## [Unreleased]
 
+### Added
+
+- Arity has been added to DefNode, BlockNode and Params. The method returns a range where the lower bound is the minimum and the upper bound is the maximum number of arguments that can be used to invoke that block/method definition.
+
 ## [5.2.0] - 2023-01-04
 
 ### Added

lib/syntax_tree/node.rb

@@ -4175,6 +4175,17 @@ def ===(other)
     def endless?
       !bodystmt.is_a?(BodyStmt)
     end
+
+    def arity
+      case params
+      when Params
+        params.arity
+      when Paren
+        params.contents.arity
+      else
+        0..0
+      end
+    end
   end
 
   # Defined represents the use of the +defined?+ operator. It can be used with
@@ -4362,6 +4373,15 @@ def keywords?
       opening.is_a?(Kw)
     end
 
+    def arity
+      case block_var
+      when BlockVar
+        block_var.params.arity
+      else
+        0..0
+      end
+    end
+
     private
 
     # If this is nested anywhere inside certain nodes, then we can't change
@@ -8325,6 +8345,23 @@ def ===(other)
         keyword_rest === other.keyword_rest && block === other.block
     end
 
+    # Returns a range representing the possible number of arguments accepted
+    # by this params node not including the block. For example:
+    #   def foo(a, b = 1, c:, d: 2, &block)
+    #     ...
+    #   end
+    # has arity 2..4
+    def arity
+      optional_keywords = keywords.count { |_label, value| value }
+      lower_bound =
+        requireds.length + posts.length + keywords.length - optional_keywords
+
+      upper_bound =
+        lower_bound + optionals.length +
+          optional_keywords if keyword_rest.nil? && rest.nil?
+      lower_bound..upper_bound
+    end
+
     private
 
     def format_contents(q, parts)

test/node_test.rb

@@ -1058,6 +1058,169 @@ def test_root_class_raises_not_implemented_errors
       end
     end
 
+    def test_arity_no_args
+      source = <<~SOURCE
+        def foo
+        end
+      SOURCE
+
+      at = location(chars: 0..11, columns: 0..3, lines: 1..2)
+      assert_node(DefNode, source, at: at) do |node|
+        assert_equal(0..0, node.arity)
+        node
+      end
+    end
+
+    def test_arity_positionals
+      source = <<~SOURCE
+        def foo(a, b = 1)
+        end
+      SOURCE
+
+      at = location(chars: 0..21, columns: 0..3, lines: 1..2)
+      assert_node(DefNode, source, at: at) do |node|
+        assert_equal(1..2, node.arity)
+        node
+      end
+    end
+
+    def test_arity_rest
+      source = <<~SOURCE
+        def foo(a, *b)
+        end
+      SOURCE
+
+      at = location(chars: 0..18, columns: 0..3, lines: 1..2)
+      assert_node(DefNode, source, at: at) do |node|
+        assert_equal(1.., node.arity)
+        node
+      end
+    end
+
+    def test_arity_keyword_rest
+      source = <<~SOURCE
+        def foo(a, **b)
+        end
+      SOURCE
+
+      at = location(chars: 0..19, columns: 0..3, lines: 1..2)
+      assert_node(DefNode, source, at: at) do |node|
+        assert_equal(1.., node.arity)
+        node
+      end
+    end
+
+    def test_arity_keywords
+      source = <<~SOURCE
+        def foo(a:, b: 1)
+        end
+      SOURCE
+
+      at = location(chars: 0..21, columns: 0..3, lines: 1..2)
+      assert_node(DefNode, source, at: at) do |node|
+        assert_equal(1..2, node.arity)
+        node
+      end
+    end
+
+    def test_arity_mixed
+      source = <<~SOURCE
+        def foo(a, b = 1, c:, d: 2)
+        end
+      SOURCE
+
+      at = location(chars: 0..31, columns: 0..3, lines: 1..2)
+      assert_node(DefNode, source, at: at) do |node|
+        assert_equal(2..4, node.arity)
+        node
+      end
+    end
+
+    guard_version("2.7.3") do
+      def test_arity_arg_forward
+        source = <<~SOURCE
+          def foo(...)
+          end
+        SOURCE
+
+        at = location(chars: 0..16, columns: 0..3, lines: 1..2)
+        assert_node(DefNode, source, at: at) do |node|
+          assert_equal(0.., node.arity)
+          node
+        end
+      end
+    end
+
+    guard_version("3.0.0") do
+      def test_arity_positional_and_arg_forward
+        source = <<~SOURCE
+          def foo(a, ...)
+          end
+        SOURCE
+
+        at = location(chars: 0..19, columns: 0..3, lines: 1..2)
+        assert_node(DefNode, source, at: at) do |node|
+          assert_equal(1.., node.arity)
+          node
+        end
+      end
+    end
+
+    def test_arity_no_parenthesis
+      source = <<~SOURCE
+        def foo a, b = 1
+        end
+      SOURCE
+
+      at = location(chars: 0..20, columns: 0..3, lines: 1..2)
+      assert_node(DefNode, source, at: at) do |node|
+        assert_equal(1..2, node.arity)
+        node
+      end
+    end
+
+    def test_block_arity_positionals
+      source = <<~SOURCE
+        [].each do |a, b, c|
+        end
+      SOURCE
+
+      at = location(chars: 8..24, columns: 8..3, lines: 1..2)
+      assert_node(BlockNode, source, at: at) do |node|
+        block = node.block
+        assert_equal(3..3, block.arity)
+        block
+      end
+    end
+
+    def test_block_arity_with_optional
+      source = <<~SOURCE
+        [].each do |a, b = 1|
+        end
+      SOURCE
+
+      at = location(chars: 8..25, columns: 8..3, lines: 1..2)
+      assert_node(BlockNode, source, at: at) do |node|
+        block = node.block
+        assert_equal(1..2, block.arity)
+        block
+      end
+    end
+
+    def test_block_arity_with_optional_keyword
+      source = <<~SOURCE
+        [].each do |a, b: 2|
+        end
+      SOURCE
+
+      at = location(chars: 8..24, columns: 8..3, lines: 1..2)
+      assert_node(BlockNode, source, at: at) do |node|
+        block = node.block
+        assert_equal(1..2, block.arity)
+        block
+      end
+    end
+
     private
 
     def location(lines: 1..1, chars: 0..0, columns: 0..0)