Add arity to CallNode, VCall, CommandCall and Command

Author: vinistock

CHANGELOG.md

@@ -9,6 +9,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/) a
 ### 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.
+- Arity has been added to CallNode, Command, CommandCall and VCall nodes. The method returns the number of arguments included in the invocation. For splats, double splats or argument forwards, this method returns Float::INFINITY.
 
 ## [5.2.0] - 2023-01-04
 

lib/syntax_tree/node.rb

@@ -775,6 +775,10 @@ def ===(other)
       other.is_a?(ArgParen) && arguments === other.arguments
     end
 
+    def arity
+      arguments&.arity || 0
+    end
+
     private
 
     def trailing_comma?
@@ -848,6 +852,22 @@ def format(q)
     def ===(other)
       other.is_a?(Args) && ArrayMatch.call(parts, other.parts)
     end
+
+    def arity
+      accepts_infinite_arguments? ? Float::INFINITY : parts.length
+    end
+
+    private
+
+    def accepts_infinite_arguments?
+      parts.any? do |part|
+        part.is_a?(ArgStar) || part.is_a?(ArgsForward) ||
+          (
+            part.is_a?(BareAssocHash) &&
+              part.assocs.any? { |p| p.is_a?(AssocSplat) }
+          )
+      end
+    end
   end
 
   # ArgBlock represents using a block operator on an expression.
@@ -1008,6 +1028,10 @@ def format(q)
     def ===(other)
       other.is_a?(ArgsForward)
     end
+
+    def arity
+      Float::INFINITY
+    end
   end
 
   # ArrayLiteral represents an array literal, which can optionally contain
@@ -3068,6 +3092,10 @@ def format_contents(q)
         end
       end
     end
+
+    def arity
+      arguments&.arity || 0
+    end
   end
 
   # Case represents the beginning of a case chain.
@@ -3481,6 +3509,10 @@ def ===(other)
         arguments === other.arguments && block === other.block
     end
 
+    def arity
+      arguments.arity
+    end
+
     private
 
     def align(q, node, &block)
@@ -3646,6 +3678,10 @@ def ===(other)
         arguments === other.arguments && block === other.block
     end
 
+    def arity
+      arguments&.arity || 0
+    end
+
     private
 
     def argument_alignment(q, doc)
@@ -11631,6 +11667,10 @@ def ===(other)
     def access_control?
       @access_control ||= %w[private protected public].include?(value.value)
     end
+
+    def arity
+      0
+    end
   end
 
   # VoidStmt represents an empty lexical block of code.

test/node_test.rb

@@ -1221,6 +1221,179 @@ def test_block_arity_with_optional_keyword
       end
     end
 
+    def test_call_node_arity_positional_arguments
+      source = <<~SOURCE
+        foo(1, 2, 3)
+      SOURCE
+
+      at = location(chars: 0..12, columns: 0..3, lines: 1..1)
+      assert_node(CallNode, source, at: at) do |node|
+        assert_equal(3, node.arity)
+        node
+      end
+    end
+
+    def test_call_node_arity_keyword_arguments
+      source = <<~SOURCE
+        foo(bar, something: 123)
+      SOURCE
+
+      at = location(chars: 0..24, columns: 0..24, lines: 1..1)
+      assert_node(CallNode, source, at: at) do |node|
+        assert_equal(2, node.arity)
+        node
+      end
+    end
+
+    def test_call_node_arity_splat_arguments
+      source = <<~SOURCE
+        foo(*bar)
+      SOURCE
+
+      at = location(chars: 0..9, columns: 0..9, lines: 1..1)
+      assert_node(CallNode, source, at: at) do |node|
+        assert_equal(Float::INFINITY, node.arity)
+        node
+      end
+    end
+
+    def test_call_node_arity_keyword_rest_arguments
+      source = <<~SOURCE
+        foo(**bar)
+      SOURCE
+
+      at = location(chars: 0..10, columns: 0..10, lines: 1..1)
+      assert_node(CallNode, source, at: at) do |node|
+        assert_equal(Float::INFINITY, node.arity)
+        node
+      end
+    end
+
+    guard_version("2.7.3") do
+      def test_call_node_arity_arg_forward_arguments
+        source = <<~SOURCE
+          def foo(...)
+            bar(...)
+          end
+        SOURCE
+
+        at = location(chars: 15..23, columns: 2..10, lines: 2..2)
+        assert_node(CallNode, source, at: at) do |node|
+          call = node.bodystmt.statements.body.first
+          assert_equal(Float::INFINITY, call.arity)
+          call
+        end
+      end
+    end
+
+    def test_command_arity_positional_arguments
+      source = <<~SOURCE
+        foo 1, 2, 3
+      SOURCE
+
+      at = location(chars: 0..11, columns: 0..3, lines: 1..1)
+      assert_node(Command, source, at: at) do |node|
+        assert_equal(3, node.arity)
+        node
+      end
+    end
+
+    def test_command_arity_keyword_arguments
+      source = <<~SOURCE
+        foo bar, something: 123
+      SOURCE
+
+      at = location(chars: 0..23, columns: 0..23, lines: 1..1)
+      assert_node(Command, source, at: at) do |node|
+        assert_equal(2, node.arity)
+        node
+      end
+    end
+
+    def test_command_arity_splat_arguments
+      source = <<~SOURCE
+        foo *bar
+      SOURCE
+
+      at = location(chars: 0..8, columns: 0..8, lines: 1..1)
+      assert_node(Command, source, at: at) do |node|
+        assert_equal(Float::INFINITY, node.arity)
+        node
+      end
+    end
+
+    def test_command_arity_keyword_rest_arguments
+      source = <<~SOURCE
+        foo **bar
+      SOURCE
+
+      at = location(chars: 0..9, columns: 0..9, lines: 1..1)
+      assert_node(Command, source, at: at) do |node|
+        assert_equal(Float::INFINITY, node.arity)
+        node
+      end
+    end
+
+    def test_command_call_arity_positional_arguments
+      source = <<~SOURCE
+        object.foo 1, 2, 3
+      SOURCE
+
+      at = location(chars: 0..18, columns: 0..3, lines: 1..1)
+      assert_node(CommandCall, source, at: at) do |node|
+        assert_equal(3, node.arity)
+        node
+      end
+    end
+
+    def test_command_call_arity_keyword_arguments
+      source = <<~SOURCE
+        object.foo bar, something: 123
+      SOURCE
+
+      at = location(chars: 0..30, columns: 0..30, lines: 1..1)
+      assert_node(CommandCall, source, at: at) do |node|
+        assert_equal(2, node.arity)
+        node
+      end
+    end
+
+    def test_command_call_arity_splat_arguments
+      source = <<~SOURCE
+        object.foo *bar
+      SOURCE
+
+      at = location(chars: 0..15, columns: 0..15, lines: 1..1)
+      assert_node(CommandCall, source, at: at) do |node|
+        assert_equal(Float::INFINITY, node.arity)
+        node
+      end
+    end
+
+    def test_command_call_arity_keyword_rest_arguments
+      source = <<~SOURCE
+        object.foo **bar
+      SOURCE
+
+      at = location(chars: 0..16, columns: 0..16, lines: 1..1)
+      assert_node(CommandCall, source, at: at) do |node|
+        assert_equal(Float::INFINITY, node.arity)
+        node
+      end
+    end
+
+    def test_vcall_arity
+      source = <<~SOURCE
+        foo
+      SOURCE
+
+      at = location(chars: 0..3, columns: 0..3, lines: 1..1)
+      assert_node(VCall, source, at: at) do |node|
+        assert_equal(0, node.arity)
+        node
+      end
+    end
+
     private
 
     def location(lines: 1..1, chars: 0..0, columns: 0..0)