Merge pull request #943 from Shopify/06-08-fix-spoom-heredoc-handling

Preserve heredoc body when translating multi-line `T.let` assertions

Author: KaanOzkan

lib/spoom/sorbet/translate/sorbet_assertions_to_rbs_comments.rb

@@ -103,10 +103,12 @@ def maybe_translate_assertion(node)
           # Otherwise, replace up to the end of the node
           end_offset = comment_end_offset || node.location.end_offset
 
+          heredoc_body = heredoc_body_within_range(value, end_offset)
+
           replacement = if node.name == :bind
             "#{rbs_annotation}#{trailing_comment}"
           else
-            "#{dedent_value(node, value)} #{rbs_annotation}#{trailing_comment}"
+            "#{dedent_value(node, value)} #{rbs_annotation}#{trailing_comment}#{heredoc_body}"
           end
 
           @rewriter << Source::Replace.new(start_offset, end_offset - 1, replacement)
@@ -212,6 +214,53 @@ def extract_trailing_comment(node)
           [" #{range.pack("C*")}", end_offset]
         end
 
+        #: (Prism::Node, Integer) -> String?
+        def heredoc_body_within_range(node, replace_end_offset)
+          heredoc_end = heredoc_end_offsets(node)
+            .select { |offset| offset <= replace_end_offset }
+            .max
+          return unless heredoc_end
+
+          opener_line_end = adjust_to_line_end(node.location.end_offset)
+          body_bytes = @ruby_bytes[(opener_line_end + 1)...heredoc_end] #: as !nil
+          body = body_bytes.pack("C*")
+          body.chomp! if @ruby_bytes[replace_end_offset] == LINE_BREAK
+          "\n#{body}"
+        end
+
+        #: (Prism::Node) -> Array[Integer]
+        def heredoc_end_offsets(node)
+          offsets = [] #: Array[Integer]
+
+          case node
+          when Prism::StringNode, Prism::InterpolatedStringNode, Prism::XStringNode, Prism::InterpolatedXStringNode
+            opening = node.opening_loc
+            closing = node.closing_loc
+            if opening && closing && opening.start_line != closing.start_line && opening.slice.start_with?("<<")
+              offsets << closing.end_offset
+            end
+          end
+
+          node.child_nodes.each do |child|
+            next unless child
+
+            offsets.concat(heredoc_end_offsets(child))
+          end
+
+          offsets
+        end
+
+        #: (Prism::Node) -> bool
+        def string_literal?(node)
+          case node
+          when Prism::StringNode, Prism::InterpolatedStringNode,
+               Prism::XStringNode, Prism::InterpolatedXStringNode
+            true
+          else
+            false
+          end
+        end
+
         #: (Prism::Node, Prism::Node) -> String
         def dedent_value(assign, value)
           if value.location.start_line == assign.location.start_line
@@ -243,7 +292,7 @@ def dedent_value(assign, value)
           # ```
           indent = value.location.start_line - assign.location.start_line
           lines = value.slice.lines
-          if lines.size > 1
+          if lines.size > 1 && !string_literal?(value)
             lines[1..]&.each_with_index do |line, i|
               lines[i + 1] = line.delete_prefix("  " * indent)
             end

rbi/spoom.rbi

@@ -3265,9 +3265,18 @@ class Spoom::Sorbet::Translate::SorbetAssertionsToRBSComments < ::Spoom::Sorbet:
   sig { params(node: ::Prism::Node).returns(T::Boolean) }
   def has_rbs_annotation?(node); end
 
+  sig { params(node: ::Prism::Node, replace_end_offset: ::Integer).returns(T.nilable(::String)) }
+  def heredoc_body_within_range(node, replace_end_offset); end
+
+  sig { params(node: ::Prism::Node).returns(T::Array[::Integer]) }
+  def heredoc_end_offsets(node); end
+
   sig { params(node: ::Prism::Node).returns(T::Boolean) }
   def maybe_translate_assertion(node); end
 
+  sig { params(node: ::Prism::Node).returns(T::Boolean) }
+  def string_literal?(node); end
+
   sig { params(node: T.nilable(::Prism::Node)).returns(T::Boolean) }
   def t?(node); end
 

test/spoom/sorbet/translate/sorbet_assertions_to_rbs_comments_test.rb

@@ -264,6 +264,93 @@ def test_translate_assigns_ignore_heredoc_values
           RB
         end
 
+        def test_translate_assigns_multiline_tlet_with_heredoc_values
+          rb = <<~RB
+            MSG = T.let(
+              <<~MSG.gsub(/[[:space:]]+/, " ").strip,
+                Do not use foo directly. Use bar instead.
+                See this guide: https://example.com/docs
+              MSG
+              String,
+            )
+
+            QUERY = T.let(
+              <<~SQL.squish.freeze,
+                SELECT id, name
+                FROM users
+                WHERE active = true
+              SQL
+              String,
+            )
+          RB
+
+          assert_equal(<<~RB, rbi_to_rbs(rb))
+            MSG = <<~MSG.gsub(/[[:space:]]+/, " ").strip #: String
+                Do not use foo directly. Use bar instead.
+                See this guide: https://example.com/docs
+              MSG
+
+            QUERY = <<~SQL.squish.freeze #: String
+                SELECT id, name
+                FROM users
+                WHERE active = true
+              SQL
+          RB
+        end
+
+        def test_translate_assigns_multiline_tlet_with_multiple_heredocs
+          rb = <<~RB
+            both = T.let(
+              foo(<<~A, <<~B),
+                first
+              A
+                second
+              B
+              String,
+            )
+          RB
+
+          assert_equal(<<~RB, rbi_to_rbs(rb))
+            both = foo(<<~A, <<~B) #: String
+                first
+              A
+                second
+              B
+          RB
+        end
+
+        def test_translate_assigns_multiline_string_literal
+          rb = <<~RB
+            s = T.let(
+              "first
+              second",
+              String,
+            )
+          RB
+
+          assert_equal(<<~RB, rbi_to_rbs(rb))
+            s = "first
+              second" #: String
+          RB
+        end
+
+        def test_translate_assigns_multiline_tlet_with_backtick_heredoc
+          rb = <<~RB
+            x = T.let(
+              <<~`CMD`,
+                echo hello
+              CMD
+              String,
+            )
+          RB
+
+          assert_equal(<<~RB, rbi_to_rbs(rb))
+            x = <<~`CMD` #: String
+                echo hello
+              CMD
+          RB
+        end
+
         def test_translate_assigns_does_not_match_bare_strings_has_heredoc
           rb = <<~RB
             a = T.let("<<~STR", String)