lib/rbi/parser.rb in rbi-0.0.14 vs lib/rbi/parser.rb in rbi-0.0.15

@@ -176,10 +176,11 @@
       @nodes_comments_assoc = nodes_comments_assoc
       @tree = T.let(Tree.new, Tree)
       @scopes_stack = T.let([@tree], T::Array[Tree])
       @last_node = T.let(nil, T.nilable(::AST::Node))
       @last_sigs = T.let([], T::Array[RBI::Sig])
+      @last_sigs_comments = T.let([], T::Array[Comment])
 
       separate_header_comments
     end
 
     sig { void }
@@ -206,15 +207,16 @@
         current_scope << parse_def(node)
       when :send
         node = parse_send(node)
         current_scope << node if node
       when :block
-        node = parse_block(node)
-        if node.is_a?(Sig)
-          @last_sigs << node
-        elsif node
-          current_scope << node
+        rbi_node = parse_block(node)
+        if rbi_node.is_a?(Sig)
+          @last_sigs << rbi_node
+          @last_sigs_comments.concat(node_comments(node))
+        elsif rbi_node
+          current_scope << rbi_node
         end
       else
         visit_all(node.children)
       end
 
@@ -266,20 +268,20 @@
         Method.new(
           node.children[0].to_s,
           params: node.children[1].children.map { |child| parse_param(child) },
           sigs: current_sigs,
           loc: loc,
-          comments: node_comments(node)
+          comments: current_sigs_comments + node_comments(node)
         )
       when :defs
         Method.new(
           node.children[1].to_s,
           params: node.children[2].children.map { |child| parse_param(child) },
           is_singleton: true,
           sigs: current_sigs,
           loc: loc,
-          comments: node_comments(node)
+          comments: current_sigs_comments + node_comments(node)
         )
       else
         raise ParseError.new("Unsupported def node type `#{node.type}`", loc)
       end
     end
@@ -322,17 +324,17 @@
       comments = node_comments(node)
 
       case method_name
       when :attr_reader
         symbols = node.children[2..-1].map { |child| child.children[0] }
-        AttrReader.new(*symbols, sigs: current_sigs, loc: loc, comments: comments)
+        AttrReader.new(*symbols, sigs: current_sigs, loc: loc, comments: current_sigs_comments + comments)
       when :attr_writer
         symbols = node.children[2..-1].map { |child| child.children[0] }
-        AttrWriter.new(*symbols, sigs: current_sigs, loc: loc, comments: comments)
+        AttrWriter.new(*symbols, sigs: current_sigs, loc: loc, comments: current_sigs_comments + comments)
       when :attr_accessor
         symbols = node.children[2..-1].map { |child| child.children[0] }
-        AttrAccessor.new(*symbols, sigs: current_sigs, loc: loc, comments: comments)
+        AttrAccessor.new(*symbols, sigs: current_sigs, loc: loc, comments: current_sigs_comments + comments)
       when :include
         names = node.children[2..-1].map { |child| parse_expr(child) }
         Include.new(*names, loc: loc, comments: comments)
       when :extend
         names = node.children[2..-1].map { |child| parse_expr(child) }
@@ -529,9 +531,16 @@
     sig { returns(T::Array[Sig]) }
     def current_sigs
       sigs = @last_sigs.dup
       @last_sigs.clear
       sigs
+    end
+
+    sig { returns(T::Array[Comment]) }
+    def current_sigs_comments
+      comments = @last_sigs_comments.dup
+      @last_sigs_comments.clear
+      comments
     end
 
     sig { void }
     def assoc_dangling_comments
       last_line = T.let(nil, T.nilable(Integer))