Merge pull request #408 from Shopify/at-fix-heredoc

Fix constants with heredoc values

Author: Morriar

lib/rbi/parser.rb

@@ -138,6 +138,17 @@ def node_string(node)
       def node_string!(node)
         T.must(node_string(node))
       end
+
+      #: (Prism::Node node) -> Prism::Location
+      def adjust_prism_location_for_heredoc(node)
+        visitor = HeredocLocationVisitor.new(
+          node.location.send(:source),
+          node.location.start_offset,
+          node.location.end_offset,
+        )
+        visitor.visit(node)
+        visitor.location
+      end
     end
 
     class TreeBuilder < Visitor
@@ -217,30 +228,40 @@ def visit_constant_assign(node)
 
         current_scope << if struct
           struct
-        elsif type_variable_definition?(node.value)
-          TypeMember.new(
-            case node
-            when Prism::ConstantWriteNode
-              node.name.to_s
-            when Prism::ConstantPathWriteNode
-              node_string!(node.target)
-            end,
-            node_string!(node.value),
-            loc: node_loc(node),
-            comments: node_comments(node),
-          )
         else
-          Const.new(
-            case node
-            when Prism::ConstantWriteNode
-              node.name.to_s
-            when Prism::ConstantPathWriteNode
-              node_string!(node.target)
-            end,
-            node_string!(node.value),
-            loc: node_loc(node),
-            comments: node_comments(node),
+          adjusted_node_location = adjust_prism_location_for_heredoc(node)
+
+          adjusted_value_location = Prism::Location.new(
+            node.value.location.send(:source),
+            node.value.location.start_offset,
+            adjusted_node_location.end_offset - node.value.location.start_offset,
           )
+
+          if type_variable_definition?(node.value)
+            TypeMember.new(
+              case node
+              when Prism::ConstantWriteNode
+                node.name.to_s
+              when Prism::ConstantPathWriteNode
+                node_string!(node.target)
+              end,
+              adjusted_value_location.slice,
+              loc: Loc.from_prism(@file, adjusted_node_location),
+              comments: node_comments(node),
+            )
+          else
+            Const.new(
+              case node
+              when Prism::ConstantWriteNode
+                node.name.to_s
+              when Prism::ConstantPathWriteNode
+                node_string!(node.target)
+              end,
+              adjusted_value_location.slice,
+              loc: Loc.from_prism(@file, adjusted_node_location),
+              comments: node_comments(node),
+            )
+          end
         end
       end
 
@@ -903,5 +924,57 @@ def visit_assoc_node(node)
         )
       end
     end
+
+    class HeredocLocationVisitor < Prism::Visitor
+      #: (Prism::Source source, Integer begin_offset, Integer end_offset) -> void
+      def initialize(source, begin_offset, end_offset)
+        super()
+        @source = source
+        @begin_offset = begin_offset
+        @end_offset = end_offset
+        @offset_last_newline = false #: bool
+      end
+
+      #: (Prism::StringNode node) -> void
+      def visit_string_node(node)
+        return unless node.heredoc?
+
+        closing_loc = node.closing_loc
+        return unless closing_loc
+
+        handle_string_node(node)
+      end
+
+      #: (Prism::InterpolatedStringNode node) -> void
+      def visit_interpolated_string_node(node)
+        return super unless node.heredoc?
+
+        closing_loc = node.closing_loc
+        return super unless closing_loc
+
+        handle_string_node(node)
+      end
+
+      #: -> Prism::Location
+      def location
+        Prism::Location.new(
+          @source,
+          @begin_offset,
+          @end_offset - @begin_offset - (@offset_last_newline ? 1 : 0),
+        )
+      end
+
+      private
+
+      #: (Prism::StringNode | Prism::InterpolatedStringNode node) -> void
+      def handle_string_node(node)
+        closing_loc = T.must(node.closing_loc)
+
+        if closing_loc.end_offset > @end_offset
+          @end_offset = closing_loc.end_offset
+          @offset_last_newline = true if node.closing_loc&.slice&.end_with?("\n")
+        end
+      end
+    end
   end
 end

lib/rbi/printer.rb

@@ -675,6 +675,13 @@ def oneline?(node)
         node.comments.empty? && node.empty?
       when Attr
         node.comments.empty? && node.sigs.empty?
+      when Const
+        return false unless node.comments.empty?
+
+        loc = node.loc
+        return true unless loc
+
+        loc.begin_line == loc.end_line
       when Method
         node.comments.empty? && node.sigs.empty? && node.params.all? { |p| p.comments.empty? }
       when Sig

rbi/rbi.rbi

@@ -1005,6 +1005,25 @@ class RBI::Parser
   end
 end
 
+class RBI::Parser::HeredocLocationVisitor < ::Prism::Visitor
+  sig { params(source: ::Prism::Source, begin_offset: ::Integer, end_offset: ::Integer).void }
+  def initialize(source, begin_offset, end_offset); end
+
+  sig { returns(::Prism::Location) }
+  def location; end
+
+  sig { params(node: ::Prism::InterpolatedStringNode).void }
+  def visit_interpolated_string_node(node); end
+
+  sig { params(node: ::Prism::StringNode).void }
+  def visit_string_node(node); end
+
+  private
+
+  sig { params(node: T.any(::Prism::InterpolatedStringNode, ::Prism::StringNode)).void }
+  def handle_string_node(node); end
+end
+
 class RBI::Parser::SigBuilder < ::RBI::Parser::Visitor
   sig { params(content: ::String, file: ::String).void }
   def initialize(content, file:); end
@@ -1119,6 +1138,9 @@ class RBI::Parser::Visitor < ::Prism::Visitor
 
   private
 
+  sig { params(node: ::Prism::Node).returns(::Prism::Location) }
+  def adjust_prism_location_for_heredoc(node); end
+
   sig { params(node: ::Prism::Node).returns(::RBI::Loc) }
   def node_loc(node); end
 

test/rbi/parser_test.rb

@@ -88,6 +88,65 @@ def test_parse_constants
       assert_equal(rbi, tree.string)
     end
 
+    def test_parse_constants_with_heredoc
+      rbi = <<~RBI
+        A = <<-EOF
+          foo
+          foo
+        EOF
+
+        B = <<~EOF.strip
+          bar
+        EOF
+
+        C = <<-'EOF'
+          baz
+        EOF
+
+        D = %Q(
+
+          qux
+
+        )
+
+        E = T.let(<<~EOF, String)
+          foo
+        EOF
+
+        F = "foo" \\
+           "bar"
+
+        G = "\#{<<~begin}\#{<<~end}"
+        begin
+          foo
+          foo
+        end
+
+        H = foo(<<~begin, <<~end)
+        begin
+          foo
+          foo
+        end
+      RBI
+
+      tree = parse_rbi(rbi)
+      assert_equal(rbi, tree.string)
+    end
+
+    def test_parse_constants_multiline_calls
+      rbi = <<~RBI
+        A = foo(
+          <<~EOF,
+            bar
+          EOF
+          baz,
+        )
+      RBI
+
+      tree = parse_rbi(rbi)
+      assert_equal(rbi, tree.string)
+    end
+
     def test_parse_constants_with_newlines
       rbi = <<~RBI
         sig { returns(Foo::

test/rbi/printer_test.rb

@@ -1113,8 +1113,10 @@ class << self; end
         class TE < T::Enum; end
         # file.rbi:1:3-2:4
         class TS < T::Struct; end
+
         # file.rbi:1:3-2:4
         C = 42
+
         # file.rbi:1:3-2:4
         extend E
         # file.rbi:1:3-2:4