Preserve sig param comments in RBS output

Author: KaanOzkan

lib/rbi/parser.rb

@@ -161,12 +161,13 @@ def t_sig_without_runtime?(node)
         !!(node.is_a?(Prism::ConstantPathNode) && node_string(node) =~ /(::)?T::Sig::WithoutRuntime/)
       end
 
-      #: (Prism::Node node) -> Array[Comment]
-      def node_comments(node)
+      #: (Prism::Node node, ?min_line: Integer?) -> Array[Comment]
+      def node_comments(node, min_line: nil)
         comments = []
 
         start_line = node.location.start_line
         start_line -= 1 unless @comments_by_line.key?(start_line)
+        start_line = [start_line, min_line].max if min_line
 
         rbs_continuation = [] #: Array[Prism::Comment]
 
@@ -767,7 +768,7 @@ def parse_params(node)
 
       #: (Prism::CallNode node) -> Sig
       def parse_sig(node)
-        builder = SigBuilder.new(@source, file: @file)
+        builder = SigBuilder.new(@source, comments_by_line: @comments_by_line, file: @file)
         builder.current.loc = node_loc(node)
         builder.visit_call_node(node)
         builder.current.comments = node_comments(node)
@@ -920,11 +921,15 @@ class SigBuilder < Visitor
       #: Sig
       attr_reader :current
 
-      #: (String content, file: String) -> void
-      def initialize(content, file:)
-        super
+      #: Integer?
+      attr_reader :params_start_line
+
+      #: (String content, comments_by_line: Hash[Integer, Prism::Comment], file: String) -> void
+      def initialize(content, comments_by_line:, file:)
+        super(content, comments_by_line: comments_by_line, file: file)
 
         @current = Sig.new #: Sig
+        @params_start_line = nil #: Integer?
       end
 
       # @override
@@ -955,7 +960,9 @@ def visit_call_node(node)
         when "overridable"
           @current.is_overridable = true
         when "params"
+          @params_start_line = node.location.start_line
           visit(node.arguments)
+          @params_start_line = nil
         when "returns"
           args = node.arguments
           if args.is_a?(Prism::ArgumentsNode)
@@ -983,6 +990,8 @@ def visit_assoc_node(node)
         @current.params << SigParam.new(
           node_string!(node.key).delete_suffix(":"),
           node_string!(node.value),
+          loc: node_loc(node),
+          comments: node_comments(node, min_line: params_start_line),
         )
       end
 

lib/rbi/rbs_printer.rb

@@ -406,7 +406,7 @@ def print_method_sig(node, sig)
       @out = old_out
 
       max_line_length = @max_line_length
-      if max_line_length.nil? || new_out.string.size <= max_line_length
+      if !sig_params_have_comments?(node, sig) && (max_line_length.nil? || new_out.string.size <= max_line_length)
         print(new_out.string)
       else
         print_method_sig_multiline(node, sig)
@@ -495,6 +495,7 @@ def print_method_sig_multiline(node, sig)
         printl("(")
         indent
         sig_params.each_with_index do |param, index|
+          print_sig_param_comments(param)
           printt
           print_sig_param(node, param)
           print(",") if index < sig_params.size - 1
@@ -899,6 +900,23 @@ def print_sig_param(node, param)
       end
     end
 
+    #: (Method node, Sig sig) -> bool
+    def sig_params_have_comments?(node, sig)
+      sig_params = sig.params
+
+      block_param = node.params.find { |param| param.is_a?(BlockParam) }
+      if block_param
+        sig_params = sig_params.reject { |param| param.name == block_param.name }
+      end
+
+      sig_params.any?(&:comments?)
+    end
+
+    #: (SigParam param) -> void
+    def print_sig_param_comments(param)
+      visit_all(param.comments)
+    end
+
     #: (Param node, last: bool) -> void
     def print_param_comment_leading_space(node, last:)
       printn

test/rbi/parser_test.rb

@@ -318,6 +318,33 @@ def test_parse_sig_comments
       assert_equal(rbi, out.string)
     end
 
+    def test_parse_sig_param_comments
+      rbi = <<~RBI
+        sig do
+          params(
+            # `a` comment
+            a: Integer,
+            # `b` comment 1
+            # `b` comment 2
+            b: String
+          ).void
+        end
+        def foo(a, b); end
+      RBI
+
+      out = Parser.parse_string(rbi)
+      assert_equal(<<~RBI, out.string)
+        sig do
+          params(
+            a: Integer, # `a` comment
+            b: String # `b` comment 1
+                      # `b` comment 2
+          ).void
+        end
+        def foo(a, b); end
+      RBI
+    end
+
     def test_parse_methods_with_visibility
       rbi = <<~RBI
         private def m1; end

test/rbi/rbs_printer_test.rb

@@ -560,6 +560,84 @@ def test_print_breaks_long_signatures
       RBI
     end
 
+    def test_print_breaks_signatures_with_sig_param_comments
+      rbi_def = Method.new("foo") do |node|
+        node.params << ReqParam.new("a")
+        node.params << ReqParam.new("b")
+      end
+
+      rbi_sig = Sig.new do |sig|
+        sig.params << SigParam.new("a", "Integer", comments: [Comment.new("First param")])
+        sig.params << SigParam.new("b", "String")
+        sig.return_type = "void"
+      end
+
+      out = StringIO.new
+      printer = RBI::RBSPrinter.new(out: out)
+      printer.print_method_sig(rbi_def, rbi_sig)
+
+      assert_equal(<<~RBI.strip, out.string)
+        (
+          # First param
+          Integer a,
+          String b
+        ) -> void
+      RBI
+    end
+
+    def test_print_breaks_signatures_with_multiple_sig_param_comments
+      rbi_def = Method.new("foo") do |node|
+        node.params << ReqParam.new("a")
+        node.params << ReqParam.new("b")
+        node.params << KwParam.new("c")
+      end
+
+      rbi_sig = Sig.new do |sig|
+        sig.params << SigParam.new("a", "Integer", comments: [Comment.new("First param")])
+        sig.params << SigParam.new("b", "String", comments: [Comment.new("Second param")])
+        sig.params << SigParam.new("c", "Symbol", comments: [Comment.new("Keyword param")])
+        sig.return_type = "void"
+      end
+
+      out = StringIO.new
+      printer = RBI::RBSPrinter.new(out: out)
+      printer.print_method_sig(rbi_def, rbi_sig)
+
+      assert_equal(<<~RBI.strip, out.string)
+        (
+          # First param
+          Integer a,
+          # Second param
+          String b,
+          # Keyword param
+          c: Symbol
+        ) -> void
+      RBI
+    end
+
+    def test_print_breaks_signatures_with_mixed_sig_param_comments
+      rbi = parse_rbi(<<~RBI)
+        sig do
+          params(
+            a: Integer,
+            # Commented param
+            b: String,
+            c: Symbol
+          ).void
+        end
+        def foo(a, b, c:); end
+      RBI
+
+      assert_equal(<<~RBI, rbi.rbs_string)
+        def foo: (
+          Integer a,
+          # Commented param
+          String b,
+          c: Symbol
+        ) -> void
+      RBI
+    end
+
     def test_print_simplified_types
       rbi = parse_rbi(<<~RBI)
         sig { returns(T.any(String, String, NilClass, T.nilable(T.nilable(Integer)), TrueClass, FalseClass)) }