Merge pull request #354 from Shopify/at-fix-parse-tenum-consts

Make TEnumBlock behave as a scope

lib/rbi/model.rb

@@ -1343,39 +1343,29 @@ def initialize(name, loc: nil, comments: [], &block)
     end
   end
 
-  class TEnumBlock < NodeWithComments
+  class TEnumBlock < Scope
     extend T::Sig
 
-    sig { returns(T::Array[String]) }
-    attr_reader :names
-
     sig do
       params(
-        names: T::Array[String],
         loc: T.nilable(Loc),
         comments: T::Array[Comment],
         block: T.nilable(T.proc.params(node: TEnumBlock).void),
       ).void
     end
-    def initialize(names = [], loc: nil, comments: [], &block)
-      super(loc: loc, comments: comments)
-      @names = names
+    def initialize(loc: nil, comments: [], &block)
+      super(loc: loc, comments: comments) {}
       block&.call(self)
     end
 
-    sig { returns(T::Boolean) }
-    def empty?
-      names.empty?
-    end
-
-    sig { params(name: String).void }
-    def <<(name)
-      @names << name
+    sig { override.returns(String) }
+    def fully_qualified_name
+      "#{parent_scope&.fully_qualified_name}.enums"
     end
 
     sig { override.returns(String) }
     def to_s
-      "#{parent_scope&.fully_qualified_name}.enums"
+      fully_qualified_name
     end
   end
 

lib/rbi/parser.rb

@@ -383,25 +383,20 @@ def visit_call_node(node)
             comments: comments,
           )
         when "enums"
-          block = node.block
-
-          unless block.is_a?(Prism::BlockNode)
-            @last_node = nil
-            return
-          end
-
-          body = block.body
-
-          unless body.is_a?(Prism::StatementsNode)
-            @last_node = nil
-            return
+          if node.block && node.arguments.nil?
+            scope = TEnumBlock.new(loc: node_loc(node), comments: node_comments(node))
+            current_scope << scope
+            @scopes_stack << scope
+            visit(node.block)
+            @scopes_stack.pop
+          else
+            current_scope << Send.new(
+              message,
+              parse_send_args(node.arguments),
+              loc: node_loc(node),
+              comments: node_comments(node),
+            )
           end
-
-          current_scope << TEnumBlock.new(
-            body.body.map { |stmt| T.cast(stmt, Prism::ConstantWriteNode).name.to_s },
-            loc: node_loc(node),
-            comments: node_comments(node),
-          )
         when "extend"
           args = node.arguments
 

lib/rbi/printer.rb

@@ -493,9 +493,7 @@ def visit_tenum_block(node)
 
       printl("enums do")
       indent
-      node.names.each do |name|
-        printl("#{name} = new")
-      end
+      visit_all(node.nodes)
       dedent
       printl("end")
     end

lib/rbi/rewriters/merge_trees.rb

@@ -361,6 +361,8 @@ def dup_empty
         Module.new(name, loc: loc, comments: comments)
       when TEnum
         TEnum.new(name, loc: loc, comments: comments)
+      when TEnumBlock
+        TEnumBlock.new(loc: loc, comments: comments)
       when TStruct
         TStruct.new(name, loc: loc, comments: comments)
       when Class
@@ -555,20 +557,6 @@ def compatible_with?(other)
     end
   end
 
-  class TEnumBlock
-    extend T::Sig
-
-    sig { override.params(other: Node).void }
-    def merge_with(other)
-      return unless other.is_a?(TEnumBlock)
-
-      super
-      other.names.each do |name|
-        names << name unless names.include?(name)
-      end
-    end
-  end
-
   class TStructProp
     extend T::Sig
 

lib/rbi/visitor.rb

@@ -39,6 +39,8 @@ def visit(node)
         visit_conflict_tree(node)
       when ScopeConflict
         visit_scope_conflict(node)
+      when TEnumBlock
+        visit_tenum_block(node)
       when Tree
         visit_tree(node)
       when Const
@@ -89,8 +91,6 @@ def visit(node)
         visit_tstruct_const(node)
       when TStructProp
         visit_tstruct_prop(node)
-      when TEnumBlock
-        visit_tenum_block(node)
       when Helper
         visit_helper(node)
       when TypeMember

test/rbi/index_test.rb

@@ -100,6 +100,8 @@ class B < T::Enum
         ::A#b=: -:3:2-3:17
         ::B: -:6:0-11:3
         ::B.enums: -:7:2-10:5
+        ::B.enums::B1: -:8:4-8:12
+        ::B.enums::B2: -:9:4-9:12
         ::C: -:15:0-15:15
         ::D: -:16:0-16:17
       IDX

test/rbi/model_test.rb

@@ -106,8 +106,10 @@ def test_model_block_builders
             end
           end
           tree << TEnum.new("Enum") do |node|
-            node << TEnumBlock.new(["A", "B"]) do |block|
+            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")])
             end
           end
           tree << Helper.new("foo") do |node|
@@ -185,7 +187,10 @@ class Struct < T::Struct
         class Enum < T::Enum
           # comment
           enums do
+            # comment
             A = new
+
+            # comment
             B = new
           end
         end
@@ -392,7 +397,9 @@ def test_model_nodes_as_strings
       mod << enum
       assert_equal("::Foo::Enum", enum.to_s)
 
-      block = TEnumBlock.new(["A", "B"])
+      block = TEnumBlock.new
+      block << Const.new("A", "new")
+      block << Const.new("B", "new")
       enum << block
       assert_equal("::Foo::Enum.enums", block.to_s)
 

test/rbi/parser_test.rb

@@ -290,14 +290,16 @@ class Foo < T::Enum
             B = new
             C = new
           end
+
           def baz; end
         end
       RBI
 
       tree = parse_rbi(rbi)
 
       # Make sure the enums are not parsed as normal classes
-      assert_equal(TEnum, tree.nodes.first.class)
+      enum = tree.nodes.first
+      assert_equal(TEnum, enum.class)
 
       assert_equal(rbi, tree.string)
     end
@@ -308,11 +310,37 @@ class Foo < T::Enum
           enums do
             A = new
           end
+
           def baz; end
         end
       RBI
 
       tree = parse_rbi(rbi)
+
+      # Make sure the enums are not parsed as normal classes
+      enum = tree.nodes.first
+      assert_equal(TEnum, enum.class)
+
+      assert_equal(rbi, tree.string)
+    end
+
+    def test_parse_t_enums_ignore_malformed_blocks
+      rbi = <<~RBI
+        class Foo < T::Enum
+          enums
+        end
+
+        class Bar < T::Enum
+          enums 1, 2
+        end
+      RBI
+
+      tree = parse_rbi(rbi)
+
+      # Make sure the enums are not parsed as normal classes
+      enum = tree.nodes.first
+      assert_equal(TEnum, enum.class)
+
       assert_equal(rbi, tree.string)
     end
 
@@ -334,6 +362,30 @@ def test_parse_helpers_vs_sends
       assert_equal(rbi, tree.string)
     end
 
+    def test_parse_t_enums_with_comments
+      rbi = <<~RBI
+        # Comment 1
+        class Foo < T::Enum
+          enums do
+            # Comment 2
+            A = new
+
+            # Comment 3
+            B = new
+
+            # Comment 4
+            C = new
+          end
+
+          # Comment 5
+          def baz; end
+        end
+      RBI
+
+      tree = Parser.parse_string(rbi)
+      assert_equal(rbi, tree.string)
+    end
+
     def test_parse_sorbet_helpers
       rbi = <<~RBI
         class Foo
@@ -657,10 +709,14 @@ def baz; end
         class Foo < T::Enum
           # -:2:2-6:5
           enums do
+            # -:3:4-3:11
             A = new
+            # -:4:4-4:11
             B = new
+            # -:5:4-5:11
             C = new
           end
+
           # -:7:2-7:14
           def baz; end
         end

test/rbi/printer_test.rb

@@ -294,8 +294,11 @@ def foo; end
 
     def test_print_t_enums
       rbi = TEnum.new("Foo")
-      block = TEnumBlock.new(["A", "B"])
-      block.names << "C"
+      block = TEnumBlock.new
+      block << Const.new("A", "new")
+      block << Const.new("B", "new")
+      block << Const.new("C", "new")
+      block << Method.new("bar")
       rbi << block
       rbi << Method.new("baz")
 
@@ -305,7 +308,9 @@ class Foo < T::Enum
             A = new
             B = new
             C = new
+            def bar; end
           end
+
           def baz; end
         end
       RBI
@@ -407,7 +412,11 @@ def test_print_nodes_with_comments
       rbi << struct
 
       enum = TEnum.new("Foo", comments: comments_multi)
-      enum << TEnumBlock.new(["A", "B"], comments: comments_single)
+      enum << TEnumBlock.new(comments: comments_single) do |block|
+        block << Const.new("A", "new")
+        block << Const.new("B", "new")
+      end
+
       rbi << enum
 
       rbi << Helper.new("foo", comments: comments_multi)