Fix parsing of dangling sigs

Signed-off-by: Alexandre Terrasa <[email protected]>
Shopify · Aug 5, 2024 · 0fd4b17 · 0fd4b17
1 parent 8e40e46
commit 0fd4b17
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 1 deletion.
diff --git a/lib/rbi/parser.rb b/lib/rbi/parser.rb
@@ -185,6 +185,7 @@ def visit_class_node(node)
         current_scope << scope
         @scopes_stack << scope
         visit(node.body)
+        scope.nodes.concat(current_sigs)
         collect_dangling_comments(node)
         @scopes_stack.pop
         @last_node = nil
@@ -269,6 +270,7 @@ def visit_module_node(node)
         current_scope << scope
         @scopes_stack << scope
         visit(node.body)
+        scope.nodes.concat(current_sigs)
         collect_dangling_comments(node)
         @scopes_stack.pop
         @last_node = nil
@@ -278,7 +280,7 @@ def visit_module_node(node)
       def visit_program_node(node)
         @last_node = node
         super
-
+        @tree.nodes.concat(current_sigs)
         collect_orphan_comments
         separate_header_comments
         set_root_tree_loc
@@ -296,6 +298,7 @@ def visit_singleton_class_node(node)
         current_scope << scope
         @scopes_stack << scope
         visit(node.body)
+        scope.nodes.concat(current_sigs)
         collect_dangling_comments(node)
         @scopes_stack.pop
         @last_node = nil

diff --git a/test/rbi/parser_test.rb b/test/rbi/parser_test.rb
@@ -153,6 +153,36 @@ def foo; end
       RBI
     end
 
+    def test_parse_dangling_sigs
+      rbi = <<~RBI
+        class Foo
+          sig { void }
+        end
+
+        module Bar
+          class << self
+            sig { void }
+          end
+          sig { void }
+        end
+        sig { void }
+        sig { returns(A) }
+      RBI
+
+      out = Parser.parse_string(rbi)
+      assert_equal(rbi, out.string)
+    end
+
+    def test_parse_sig_standalone
+      rbi = <<~RBI
+        sig { void }
+        sig { returns(A) }
+      RBI
+
+      out = Parser.parse_string(rbi)
+      assert_equal(rbi, out.string)
+    end
+
     def test_parse_methods_with_visibility
       rbi = <<~RBI
         private def m1; end