diff --git a/lib/spoom/sorbet/translate/sorbet_assertions_to_rbs_comments.rb b/lib/spoom/sorbet/translate/sorbet_assertions_to_rbs_comments.rb index d2ba75e4..9f01aa1a 100644 --- a/lib/spoom/sorbet/translate/sorbet_assertions_to_rbs_comments.rb +++ b/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,42 @@ 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, Prism::Node) -> String def dedent_value(assign, value) if value.location.start_line == assign.location.start_line diff --git a/rbi/spoom.rbi b/rbi/spoom.rbi index ef1a2327..03c8118f 100644 --- a/rbi/spoom.rbi +++ b/rbi/spoom.rbi @@ -3043,6 +3043,12 @@ 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 diff --git a/test/spoom/sorbet/translate/sorbet_assertions_to_rbs_comments_test.rb b/test/spoom/sorbet/translate/sorbet_assertions_to_rbs_comments_test.rb index ed57d175..163944b2 100644 --- a/test/spoom/sorbet/translate/sorbet_assertions_to_rbs_comments_test.rb +++ b/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)