Parse T::Enum values as special nodes

Signed-off-by: Alexandre Terrasa <alexandre.terrasa@shopify.com>

Author: Morriar

lib/rbi/index.rb

@@ -217,7 +217,17 @@ class TEnumBlock
     # @override
     #: -> Array[String]
     def index_ids
-      [to_s]
+      [fully_qualified_name]
+    end
+  end
+
+  class TEnumValue
+    include Indexable
+
+    # @override
+    #: -> Array[String]
+    def index_ids
+      [fully_qualified_name]
     end
   end
 end

lib/rbi/model.rb

@@ -1037,6 +1037,29 @@ def to_s
     end
   end
 
+  class TEnumValue < NodeWithComments
+    #: String
+    attr_reader :name
+
+    #: (String name, ?loc: Loc?, ?comments: Array[Comment]) ?{ (TEnumValue node) -> void } -> void
+    def initialize(name, loc: nil, comments: [], &block)
+      super(loc: loc, comments: comments)
+      @name = name
+      block&.call(self)
+    end
+
+    #: -> String
+    def fully_qualified_name
+      "#{parent_scope&.fully_qualified_name}::#{name}"
+    end
+
+    # @override
+    #: -> String
+    def to_s
+      fully_qualified_name
+    end
+  end
+
   # Sorbet's misc.
 
   class Helper < NodeWithComments

lib/rbi/parser.rb

@@ -229,6 +229,17 @@ def visit_constant_assign(node)
             loc: node_loc(node),
             comments: node_comments(node),
           )
+        elsif t_enum_value?(node)
+          TEnumValue.new(
+            case node
+            when Prism::ConstantWriteNode
+              node.name.to_s
+            when Prism::ConstantPathWriteNode
+              node_string!(node.target)
+            end,
+            loc: node_loc(node),
+            comments: node_comments(node),
+          )
         else
           Const.new(
             case node
@@ -819,6 +830,11 @@ def set_root_tree_loc
       def type_variable_definition?(node)
         node.is_a?(Prism::CallNode) && (node.message == "type_member" || node.message == "type_template")
       end
+
+      #: (Prism::Node? node) -> bool
+      def t_enum_value?(node)
+        current_scope.is_a?(TEnumBlock) && node.is_a?(Prism::ConstantWriteNode) && node.value.slice == "new"
+      end
     end
 
     class SigBuilder < Visitor

lib/rbi/printer.rb

@@ -542,6 +542,16 @@ def visit_tenum_block(node)
       printl("end")
     end
 
+    # @override
+    #: (TEnumValue node) -> void
+    def visit_tenum_value(node)
+      print_blank_line_before(node)
+      print_loc(node)
+      visit_all(node.comments)
+
+      printl("#{node.name} = new")
+    end
+
     # @override
     #: (TypeMember node) -> void
     def visit_type_member(node)

lib/rbi/rbs_printer.rb

@@ -645,19 +645,22 @@ def visit_tenum(node)
     # @override
     #: (TEnumBlock node) -> void
     def visit_tenum_block(node)
-      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,
-          )
-        else
-          child
-        end
-        visit(child)
-        @previous_node = child
+      visit_all(node.nodes)
+    end
+
+    # @override
+    #: (TEnumValue node) -> void
+    def visit_tenum_value(node)
+      print_blank_line_before(node)
+      print_loc(node)
+      visit_all(node.comments)
+
+      t_enum = node.parent_scope&.parent_scope
+
+      if t_enum.is_a?(TEnum)
+        printl("#{node.name}: #{t_enum.name}")
+      else
+        printl("#{node.name}: untyped")
       end
     end
 

lib/rbi/visitor.rb

@@ -88,6 +88,8 @@ def visit(node)
         visit_sig(node)
       when SigParam
         visit_sig_param(node)
+      when TEnumValue
+        visit_tenum_value(node)
       when TStructConst
         visit_tstruct_const(node)
       when TStructProp
@@ -222,6 +224,9 @@ def visit_tenum(node); end
     #: (TEnumBlock node) -> void
     def visit_tenum_block(node); end
 
+    #: (TEnumValue node) -> void
+    def visit_tenum_value(node); end
+
     #: (Helper node) -> void
     def visit_helper(node); end
 

test/rbi/parser_test.rb

@@ -309,7 +309,14 @@ def baz; end
 
       # Make sure the enums are not parsed as normal classes
       enum = tree.nodes.first
-      assert_equal(TEnum, enum.class)
+      assert_instance_of(TEnum, enum)
+
+      block = T.cast(enum, TEnum).nodes.first
+      assert_instance_of(TEnumBlock, block)
+
+      values = T.cast(block, TEnumBlock).nodes
+      assert_equal(3, values.size)
+      assert_equal(3, values.grep(TEnumValue).size)
 
       assert_equal(rbi, tree.string)
     end
@@ -329,7 +336,13 @@ def baz; end
 
       # Make sure the enums are not parsed as normal classes
       enum = tree.nodes.first
-      assert_equal(TEnum, enum.class)
+      assert_instance_of(TEnum, enum)
+
+      block = T.cast(enum, TEnum).nodes.first
+      assert_instance_of(TEnumBlock, block)
+
+      value = T.cast(block, TEnumBlock).nodes.first
+      assert_instance_of(TEnumValue, value)
 
       assert_equal(rbi, tree.string)
     end