Handle for _ in _ procs in Prism.node_for

Author: eregon

lib/prism/parse_result.rb

@@ -946,10 +946,12 @@ def self.node_for(callable)
       when CallNode
         # Proc#source_location returns start_column 5 for `proc { ... }` (the `{`)
         node.block.is_a?(BlockNode) && node.block.opening_loc.start_offset == start_offset && node.end_offset == end_offset
+      when ForNode
+        node.start_offset == start_offset && node.end_offset == end_offset
       else
         false
       end
-    end #: DefNode | LambdaNode | CallNode
+    end #: DefNode | LambdaNode | CallNode | ForNode
 
     raise ArgumentError, "Could not find node for #{callable} in #{file} at (#{start_line},#{start_column})-(#{end_line},#{end_column})" unless found
     found

templates/sig/prism.rbs.erb

@@ -94,5 +94,5 @@ module Prism
 
   def self.scope: (?locals: Array[Symbol], ?forwarding: Array[Symbol]) -> Scope
 
-  def self.node_for: (Method | UnboundMethod | Proc callable) -> (DefNode | LambdaNode | CallNode)
+  def self.node_for: (Method | UnboundMethod | Proc callable) -> (DefNode | LambdaNode | CallNode | ForNode)
 end

test/prism/ruby/node_for_test.rb

@@ -24,6 +24,17 @@ def return_block(&block)
       block
     end
 
+    iter = Object.new
+    def iter.each(&block)
+      block.call(block)
+    end
+
+    for pr in iter
+      42
+    end
+    FOR_BODY_PROC = pr
+    FOR_BODY_PROC_LOCATION = [__LINE__-4, 4, __LINE__-2, 7]
+
     def with_location(callable, locs, file = __FILE__)
       source_location = [file, *locs]
       if RUBY_VERSION >= "4.0"
@@ -95,6 +106,13 @@ def test_method_to_proc
       assert_equal "def inline_method = 42", node.slice
     end
 
+    def test_for
+      node = Prism.node_for(with_location(FOR_BODY_PROC, FOR_BODY_PROC_LOCATION))
+      assert_instance_of(Prism::ForNode, node)
+      assert_equal "for pr in iter\n#{INDENT}  42\n#{INDENT}end", node.slice
+      assert_equal "42", node.statements.slice
+    end
+
     def test_eval
       l = with_location(eval("-> { 42 }"), [1, 2, 1, 9], "(eval at #{__FILE__}:#{__LINE__})")
       e = assert_raise(ArgumentError) { Prism.node_for(l) }