WIP

Author: Morriar

lib/rbi/model.rb

@@ -269,13 +269,16 @@ def fully_qualified_name
 
   class Const < NodeWithComments
     #: String
-    attr_reader :name, :value
+    attr_reader :name
 
-    #: (String name, String value, ?loc: Loc?, ?comments: Array[Comment]) ?{ (Const node) -> void } -> void
-    def initialize(name, value, loc: nil, comments: [], &block)
+    #: (String | Type)?
+    attr_reader :type
+
+    #: (String name, (String | Type)? type, ?loc: Loc?, ?comments: Array[Comment]) ?{ (Const node) -> void } -> void
+    def initialize(name, type: nil, loc: nil, comments: [], &block)
       super(loc: loc, comments: comments)
       @name = name
-      @value = value
+      @type = type
       block&.call(self)
     end
 

lib/rbi/parser.rb

@@ -237,7 +237,7 @@ def visit_constant_assign(node)
             when Prism::ConstantPathWriteNode
               node_string!(node.target)
             end,
-            node_string!(node.value),
+            type: parse_t_let_type(node.value),
             loc: node_loc(node),
             comments: node_comments(node),
           )
@@ -791,6 +791,24 @@ def parse_visibility(name, node)
         end
       end
 
+      # Parse a string containing a `T.let(x, X)` and extract the type
+      #
+      # Returns `nil` is the string is not a `T.let`.
+      #: (Prism::Node? node) -> String?
+      def parse_t_let_type(node)
+        return unless node
+
+        return unless node.is_a?(Prism::CallNode)
+        return unless node.name == :let
+        return unless node.receiver&.slice =~ /^(::)?T$/
+
+        arguments = node.arguments&.arguments
+        return unless arguments
+        return unless arguments.size == 2
+
+        arguments.fetch(1, nil)&.slice
+      end
+
       #: -> void
       def separate_header_comments
         current_scope.nodes.dup.each do |child_node|

lib/rbi/printer.rb

@@ -233,7 +233,11 @@ def visit_const(node)
       print_loc(node)
       visit_all(node.comments)
 
-      printl("#{node.name} = #{node.value}")
+      if node.type
+        printl("#{node.name} = T.let(T.unsafe(nil), #{node.type})")
+      else
+        printl("#{node.name} = T.unsafe(nil)")
+      end
     end
 
     # @override
@@ -537,7 +541,19 @@ def visit_tenum_block(node)
 
       printl("enums do")
       indent
-      visit_all(node.nodes)
+      @previous_node = nil
+      node.nodes.each do |child|
+        if child.is_a?(Const) && child.type.nil?
+          print_blank_line_before(child)
+          print_loc(child)
+          visit_all(child.comments)
+          printl("#{child.name} = new")
+        else
+          visit(child)
+        end
+        @previous_node = child
+      end
+      @previous_node = node
       dedent
       printl("end")
     end

lib/rbi/rbs_printer.rb

@@ -229,7 +229,7 @@ def visit_const(node)
       print_loc(node)
       visit_all(node.comments)
 
-      type = parse_t_let(node.value)
+      type = node.type
       if type
         type = parse_type(type)
         printl("#{node.name}: #{type.rbs_string}")
@@ -645,20 +645,20 @@ def visit_tenum(node)
     # @override
     #: (TEnumBlock node) -> void
     def visit_tenum_block(node)
+      @previous_node = nil
+      parent = node.parent_scope
       node.nodes.each do |child|
-        child = if child.is_a?(Const) && child.value == "new"
-          parent = node.parent_scope
-          Const.new(
-            child.name,
-            "T.let(nil, #{parent.is_a?(TEnum) ? parent.name : "T.untyped"})",
-            comments: child.comments,
-          )
+        if child.is_a?(Const) && child.type.nil?
+          print_blank_line_before(child)
+          print_loc(child)
+          visit_all(child.comments)
+          printl("#{child.name}: #{parent.is_a?(TEnum) ? parent.name : "untyped"}")
         else
-          child
+          visit(child)
         end
-        visit(child)
         @previous_node = child
       end
+      @previous_node = node
     end
 
     # @override
@@ -850,30 +850,6 @@ def parse_type(type)
     rescue Type::Error => e
       raise Error, "Failed to parse type `#{type}` (#{e.message})"
     end
-
-    # Parse a string containing a `T.let(x, X)` and extract the type
-    #
-    # Returns `nil` is the string is not a `T.let`.
-    #: (String? code) -> String?
-    def parse_t_let(code)
-      return unless code
-
-      res = Prism.parse(code)
-      return unless res.success?
-
-      node = res.value
-      return unless node.is_a?(Prism::ProgramNode)
-
-      node = node.statements.body.first
-      return unless node.is_a?(Prism::CallNode)
-      return unless node.name == :let
-      return unless node.receiver&.slice =~ /^(::)?T$/
-
-      arguments = node.arguments&.arguments
-      return unless arguments
-
-      arguments.fetch(1, nil)&.slice
-    end
   end
 
   class TypePrinter

lib/rbi/rewriters/merge_trees.rb

@@ -384,7 +384,7 @@ class Const
     # @override
     #: (Node other) -> bool
     def compatible_with?(other)
-      other.is_a?(Const) && name == other.name && value == other.value
+      other.is_a?(Const) && name == other.name && type == other.type
     end
   end
 

test/rbi/model_test.rb

@@ -42,7 +42,7 @@ def test_model_block_builders
           tree << SingletonClass.new do |node|
             node.comments << Comment.new("comment")
           end
-          tree << Const.new("C", "42") do |node|
+          tree << Const.new("C", type: "Integer") do |node|
             node.comments << Comment.new("comment")
           end
           tree << AttrAccessor.new(:a1) do |node|
@@ -108,8 +108,8 @@ def test_model_block_builders
           tree << TEnum.new("Enum") do |node|
             node << TEnumBlock.new do |block|
               block.comments << Comment.new("comment")
-              block << Const.new("A", "new", comments: [Comment.new("comment")])
-              block << Const.new("B", "new", comments: [Comment.new("comment")])
+              block << Const.new("A", comments: [Comment.new("comment")])
+              block << Const.new("B", comments: [Comment.new("comment")])
             end
           end
           tree << Helper.new("foo") do |node|
@@ -270,19 +270,19 @@ def test_model_fully_qualified_names
       cls1 << singleton_class
       assert_equal("::Foo::Bar::<self>", singleton_class.fully_qualified_name)
 
-      const = Const.new("Foo", "42")
+      const = Const.new("Foo", type: "Integer")
       assert_equal("::Foo", const.fully_qualified_name)
 
       mod << const
       assert_equal("::Foo::Foo", const.fully_qualified_name)
 
-      const2 = Const.new("Foo::Bar", "42")
+      const2 = Const.new("Foo::Bar", type: "42")
       assert_equal("::Foo::Bar", const2.fully_qualified_name)
 
       mod << const2
       assert_equal("::Foo::Foo::Bar", const2.fully_qualified_name)
 
-      const3 = Const.new("::Foo::Bar", "42")
+      const3 = Const.new("::Foo::Bar", type: "42")
       assert_equal("::Foo::Bar", const3.fully_qualified_name)
 
       mod << const3
@@ -344,13 +344,13 @@ def test_model_nodes_as_strings
       cls << singleton_class
       assert_equal("::Foo::Bar::<self>", singleton_class.to_s)
 
-      const = Const.new("Foo", "42")
+      const = Const.new("Foo", type: "Integer")
       assert_equal("::Foo", const.to_s)
 
       mod << const
       assert_equal("::Foo::Foo", const.to_s)
 
-      const2 = Const.new("Foo::Bar", "42")
+      const2 = Const.new("Foo::Bar", type: "42")
       assert_equal("::Foo::Bar", const2.to_s)
 
       mod << const2
@@ -404,8 +404,8 @@ def test_model_nodes_as_strings
       assert_equal("::Foo::Enum", enum.to_s)
 
       block = TEnumBlock.new
-      block << Const.new("A", "new")
-      block << Const.new("B", "new")
+      block << Const.new("A", comments: [Comment.new("comment")])
+      block << Const.new("B", comments: [Comment.new("comment")])
       enum << block
       assert_equal("::Foo::Enum.enums", block.to_s)
 

test/rbi/parser_test.rb

@@ -85,7 +85,16 @@ def test_parse_constants
       RBI
 
       tree = parse_rbi(rbi)
-      assert_equal(rbi, tree.string)
+      assert_equal(<<~RBI, tree.string)
+        Foo = T.unsafe(nil)
+        Bar = T.unsafe(nil)
+        Baz = T.unsafe(nil)
+        A = T.unsafe(nil)
+        B = T.unsafe(nil)
+        C = T.unsafe(nil)
+        D = T.unsafe(nil)
+        A::B::C = T.unsafe(nil)
+      RBI
     end
 
     def test_parse_constants_with_newlines
@@ -546,13 +555,13 @@ def test_parse_constants_locations
       tree = parse_rbi(rbi)
       assert_equal(<<~RBI, tree.string(print_locs: true))
         # -:1:0-1:8
-        Foo = 42
+        Foo = T.unsafe(nil)
         # -:2:0-2:11
-        Bar = "foo"
+        Bar = T.unsafe(nil)
         # -:3:0-3:11
-        ::Baz = Bar
+        ::Baz = T.unsafe(nil)
         # -:4:0-4:13
-        A::B::C = Foo
+        A::B::C = T.unsafe(nil)
       RBI
     end
 
@@ -889,7 +898,7 @@ def c(a); end
             attr_reader :a
 
             # E comment
-            E = _
+            E = T.unsafe(nil)
           end
         end
       RBI

test/rbi/printer_test.rb

@@ -96,14 +96,14 @@ def baz_class_method; end
 
     def test_print_constants
       rbi = Tree.new
-      rbi << Const.new("Foo", "42")
-      rbi << Const.new("Bar", "'foo'")
-      rbi << Const.new("Baz", "Bar")
+      rbi << Const.new("Foo", type: "Integer")
+      rbi << Const.new("Bar", type: Type.boolean)
+      rbi << Const.new("Baz")
 
       assert_equal(<<~RBI, rbi.string)
-        Foo = 42
-        Bar = 'foo'
-        Baz = Bar
+        Foo = T.let(T.unsafe(nil), Integer)
+        Bar = T.let(T.unsafe(nil), T::Boolean)
+        Baz = T.unsafe(nil)
       RBI
     end
 
@@ -300,9 +300,9 @@ def foo; end
     def test_print_t_enums
       rbi = TEnum.new("Foo")
       block = TEnumBlock.new
-      block << Const.new("A", "new")
-      block << Const.new("B", "new")
-      block << Const.new("C", "new")
+      block << Const.new("A")
+      block << Const.new("B")
+      block << Const.new("C", type: "Integer")
       block << Method.new("bar")
       rbi << block
       rbi << Method.new("baz")
@@ -312,7 +312,7 @@ class Foo < T::Enum
           enums do
             A = new
             B = new
-            C = new
+            C = T.let(T.unsafe(nil), Integer)
             def bar; end
           end
 
@@ -401,7 +401,7 @@ def test_print_nodes_with_comments
       rbi << Module.new("Foo", comments: comments_single)
       rbi << Class.new("Bar", comments: comments_multi)
       rbi << SingletonClass.new(comments: comments_single)
-      rbi << Const.new("Foo", "42", comments: comments_multi)
+      rbi << Const.new("Foo", comments: comments_multi)
       rbi << Include.new("A", comments: comments_single)
       rbi << Extend.new("A", comments: comments_multi)
 
@@ -418,8 +418,8 @@ def test_print_nodes_with_comments
 
       enum = TEnum.new("Foo", comments: comments_multi)
       enum << TEnumBlock.new(comments: comments_single) do |block|
-        block << Const.new("A", "new")
-        block << Const.new("B", "new")
+        block << Const.new("A")
+        block << Const.new("B")
       end
 
       rbi << enum
@@ -442,7 +442,7 @@ class << self; end
 
         # This is a
         # Multiline Comment
-        Foo = 42
+        Foo = T.unsafe(nil)
 
         # This is a single line comment
         include A
@@ -1092,7 +1092,7 @@ def test_print_nodes_locations
       rbi << SingletonClass.new(loc: loc)
       rbi << TEnum.new("TE", loc: loc)
       rbi << TStruct.new("TS", loc: loc)
-      rbi << Const.new("C", "42", loc: loc)
+      rbi << Const.new("C", loc: loc)
       rbi << Extend.new("E", loc: loc)
       rbi << Include.new("I", loc: loc)
       rbi << Send.new("foo", loc: loc)
@@ -1114,7 +1114,7 @@ class TE < T::Enum; end
         # file.rbi:1:3-2:4
         class TS < T::Struct; end
         # file.rbi:1:3-2:4
-        C = 42
+        C = T.unsafe(nil)
         # file.rbi:1:3-2:4
         extend E
         # file.rbi:1:3-2:4

test/rbi/rbs_printer_test.rb

@@ -709,8 +709,8 @@ def test_print_t_enums
         class Foo < T::Enum
           enums do
             A = new
-            B = new
-            C = new
+            B = 42
+            C = T.let(nil, T.nilable(Foo))
           end
 
           def baz; end
@@ -720,8 +720,8 @@ def baz; end
       assert_equal(<<~RBI, rbi.rbs_string)
         class Foo
           A: Foo
-          B: Foo
-          C: Foo
+          B: unsafe
+          C: Foo?
           def baz: -> untyped
         end
       RBI