Merge pull request #353 from Shopify/at-fix-parse-sig-comments

Fix parsing of comments on signatures

lib/rbi/model.rb

@@ -1101,7 +1101,7 @@ def to_s
 
   # Sorbet's sigs
 
-  class Sig < Node
+  class Sig < NodeWithComments
     extend T::Sig
 
     sig { returns(T::Array[SigParam]) }
@@ -1130,6 +1130,7 @@ class Sig < Node
         type_params: T::Array[String],
         checked: T.nilable(Symbol),
         loc: T.nilable(Loc),
+        comments: T::Array[Comment],
         block: T.nilable(T.proc.params(node: Sig).void),
       ).void
     end
@@ -1143,9 +1144,10 @@ def initialize(
       type_params: [],
       checked: nil,
       loc: nil,
+      comments: [],
       &block
     )
-      super(loc: loc)
+      super(loc: loc, comments: comments)
       @params = params
       @return_type = return_type
       @is_abstract = is_abstract

lib/rbi/parser.rb

@@ -169,7 +169,6 @@ def initialize(source, comments:, file:)
         @scopes_stack = T.let([@tree], T::Array[Tree])
         @last_node = T.let(nil, T.nilable(Prism::Node))
         @last_sigs = T.let([], T::Array[RBI::Sig])
-        @last_sigs_comments = T.let([], T::Array[Comment])
       end
 
       sig { override.params(node: Prism::ClassNode).void }
@@ -260,13 +259,14 @@ def visit_def_node(node)
 
         # We need to collect the comments with `current_sigs_comments` _before_ visiting the parameters to make sure
         # the method comments are properly associated with the sigs and not the parameters.
-        comments = current_sigs_comments + node_comments(node)
+        sigs = current_sigs
+        comments = detach_comments_from_sigs(sigs) + node_comments(node)
         params = parse_params(node.parameters)
 
         current_scope << Method.new(
           node.name.to_s,
           params: params,
-          sigs: current_sigs,
+          sigs: sigs,
           loc: node_loc(node),
           comments: comments,
           is_singleton: !!node.receiver,
@@ -339,11 +339,14 @@ def visit_call_node(node)
             return
           end
 
+          sigs = current_sigs
+          comments = detach_comments_from_sigs(sigs) + node_comments(node)
+
           current_scope << AttrReader.new(
             *T.unsafe(args.arguments.map { |arg| node_string!(arg).delete_prefix(":").to_sym }),
-            sigs: current_sigs,
+            sigs: sigs,
             loc: node_loc(node),
-            comments: current_sigs_comments + node_comments(node),
+            comments: comments,
           )
         when "attr_writer"
           args = node.arguments
@@ -353,11 +356,14 @@ def visit_call_node(node)
             return
           end
 
+          sigs = current_sigs
+          comments = detach_comments_from_sigs(sigs) + node_comments(node)
+
           current_scope << AttrWriter.new(
             *T.unsafe(args.arguments.map { |arg| node_string!(arg).delete_prefix(":").to_sym }),
-            sigs: current_sigs,
+            sigs: sigs,
             loc: node_loc(node),
-            comments: current_sigs_comments + node_comments(node),
+            comments: comments,
           )
         when "attr_accessor"
           args = node.arguments
@@ -367,11 +373,14 @@ def visit_call_node(node)
             return
           end
 
+          sigs = current_sigs
+          comments = detach_comments_from_sigs(sigs) + node_comments(node)
+
           current_scope << AttrAccessor.new(
             *T.unsafe(args.arguments.map { |arg| node_string!(arg).delete_prefix(":").to_sym }),
-            sigs: current_sigs,
+            sigs: sigs,
             loc: node_loc(node),
-            comments: current_sigs_comments + node_comments(node),
+            comments: comments,
           )
         when "enums"
           block = node.block
@@ -548,10 +557,15 @@ def current_sigs
         sigs
       end
 
-      sig { returns(T::Array[Comment]) }
-      def current_sigs_comments
-        comments = @last_sigs_comments.dup
-        @last_sigs_comments.clear
+      sig { params(sigs: T::Array[Sig]).returns(T::Array[Comment]) }
+      def detach_comments_from_sigs(sigs)
+        comments = T.let([], T::Array[Comment])
+
+        sigs.each do |sig|
+          comments += sig.comments.dup
+          sig.comments.clear
+        end
+
         comments
       end
 
@@ -678,11 +692,10 @@ def parse_params(node)
 
       sig { params(node: Prism::CallNode).returns(Sig) }
       def parse_sig(node)
-        @last_sigs_comments = node_comments(node)
-
         builder = SigBuilder.new(@source, file: @file)
         builder.current.loc = node_loc(node)
         builder.visit_call_node(node)
+        builder.current.comments = node_comments(node)
         builder.current
       end
 

lib/rbi/printer.rb

@@ -426,6 +426,7 @@ def visit_kw_arg(node)
     sig { override.params(node: Sig).void }
     def visit_sig(node)
       print_loc(node)
+      visit_all(node.comments)
 
       max_line_length = self.max_line_length
       if oneline?(node) && max_line_length.nil?

test/rbi/parser_test.rb

@@ -185,6 +185,19 @@ def test_parse_sig_standalone
       assert_equal(rbi, out.string)
     end
 
+    def test_parse_sig_comments
+      rbi = <<~RBI
+        # Sig comment
+        sig { void }
+        # Multi line
+        # sig comment
+        sig { void }
+      RBI
+
+      out = Parser.parse_string(rbi)
+      assert_equal(rbi, out.string)
+    end
+
     def test_parse_methods_with_visibility
       rbi = <<~RBI
         private def m1; end