Only delete extend T::Helpers when needed

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

Author: Morriar

lib/spoom/sorbet/translate/sorbet_sigs_to_rbs_comments.rb

@@ -12,11 +12,12 @@ def initialize(ruby_contents, file:, positional_names:)
           super(ruby_contents, file: file)
 
           @positional_names = positional_names #: bool
-          @nesting = [] #: Array[Prism::ClassNode | Prism::ModuleNode | Prism::SingletonClassNode]
           @last_sigs = [] #: Array[[Prism::CallNode, RBI::Sig]]
+          @class_annotations = [] #: Array[Prism::CallNode]
           @type_members = [] #: Array[String]
           @extend_t_helpers = [] #: Array[Prism::CallNode]
           @extend_t_generics = [] #: Array[Prism::CallNode]
+          @seen_mixes_in_class_methods = false #: bool
         end
 
         # @override
@@ -73,7 +74,9 @@ def visit_call_node(node)
           when "extend"
             visit_extend(node)
           when "abstract!", "interface!", "sealed!", "final!", "requires_ancestor"
-            visit_class_annotation(node)
+            @class_annotations << node
+          when "mixes_in_class_methods"
+            @seen_mixes_in_class_methods = true
           else
             super
           end
@@ -99,25 +102,39 @@ def visit_constant_write_node(node)
 
         #: (Prism::ClassNode | Prism::ModuleNode | Prism::SingletonClassNode) { -> void } -> void
         def visit_scope(node, &block)
-          @nesting << node
+          old_class_annotations = @class_annotations
+          @class_annotations = []
           old_type_members = @type_members
           @type_members = []
           old_extend_t_helpers = @extend_t_helpers
           @extend_t_helpers = []
           old_extend_t_generics = @extend_t_generics
           @extend_t_generics = []
+          old_seen_mixes_in_class_methods = @seen_mixes_in_class_methods
+          @seen_mixes_in_class_methods = false
 
           yield
 
+          delete_extend_t_generics
+
           if @type_members.any?
             indent = " " * node.location.start_column
             @rewriter << Source::Insert.new(node.location.start_offset, "#: [#{@type_members.join(", ")}]\n#{indent}")
           end
 
+          unless @seen_mixes_in_class_methods
+            delete_extend_t_helpers
+          end
+
+          @class_annotations.each do |call|
+            apply_class_annotation(node, call)
+          end
+
+          @class_annotations = old_class_annotations
           @type_members = old_type_members
           @extend_t_helpers = old_extend_t_helpers
           @extend_t_generics = old_extend_t_generics
-          @nesting.pop
+          @seen_mixes_in_class_methods = old_seen_mixes_in_class_methods
         end
 
         #: (Prism::CallNode) -> void
@@ -177,8 +194,8 @@ def visit_extend(node)
           end
         end
 
-        #: (Prism::CallNode node) -> void
-        def visit_class_annotation(node)
+        #: (Prism::ClassNode | Prism::ModuleNode | Prism::SingletonClassNode, Prism::CallNode) -> void
+        def apply_class_annotation(parent, node)
           unless node.message == "abstract!" || node.message == "interface!" || node.message == "sealed!" ||
               node.message == "final!" || node.message == "requires_ancestor"
             raise Error, "Expected abstract!, interface!, sealed!, final!, or requires_ancestor"
@@ -187,18 +204,17 @@ def visit_class_annotation(node)
           return unless node.receiver.nil? || node.receiver.is_a?(Prism::SelfNode)
           return unless node.arguments.nil?
 
-          klass = @nesting.last #: as Prism::Node
-          indent = " " * klass.location.start_column
+          indent = " " * parent.location.start_column
 
           case node.message
           when "abstract!"
-            @rewriter << Source::Insert.new(klass.location.start_offset, "# @abstract\n#{indent}")
+            @rewriter << Source::Insert.new(parent.location.start_offset, "# @abstract\n#{indent}")
           when "interface!"
-            @rewriter << Source::Insert.new(klass.location.start_offset, "# @interface\n#{indent}")
+            @rewriter << Source::Insert.new(parent.location.start_offset, "# @interface\n#{indent}")
           when "sealed!"
-            @rewriter << Source::Insert.new(klass.location.start_offset, "# @sealed\n#{indent}")
+            @rewriter << Source::Insert.new(parent.location.start_offset, "# @sealed\n#{indent}")
           when "final!"
-            @rewriter << Source::Insert.new(klass.location.start_offset, "# @final\n#{indent}")
+            @rewriter << Source::Insert.new(parent.location.start_offset, "# @final\n#{indent}")
           when "requires_ancestor"
             block = node.block
             return unless block.is_a?(Prism::BlockNode)
@@ -209,7 +225,7 @@ def visit_class_annotation(node)
 
             arg = body.body.first #: as Prism::Node
             srb_type = RBI::Type.parse_node(arg)
-            @rewriter << Source::Insert.new(klass.location.start_offset, "# @requires_ancestor: #{srb_type.rbs_string}\n#{indent}")
+            @rewriter << Source::Insert.new(parent.location.start_offset, "# @requires_ancestor: #{srb_type.rbs_string}\n#{indent}")
           end
 
           from = adjust_to_line_start(node.location.start_offset)
@@ -305,6 +321,8 @@ def delete_extend_t_helpers
             to = adjust_to_new_line(to)
             @rewriter << Source::Delete.new(from, to)
           end
+
+          @extend_t_helpers.clear
         end
 
         #: -> void
@@ -315,6 +333,8 @@ def delete_extend_t_generics
             to = adjust_to_new_line(to)
             @rewriter << Source::Delete.new(from, to)
           end
+
+          @extend_t_generics.clear
         end
       end
     end

test/spoom/sorbet/translate/sorbet_sigs_to_rbs_comments_test.rb

@@ -226,8 +226,8 @@ class << self
             # @abstract
             # @requires_ancestor: singleton(Foo::Bar)
             class A
-              # @interface
               # @sealed
+              # @interface
               module B
                 # @final
                 class << self
@@ -249,9 +249,34 @@ def bar
                 end
               end
             end
+
+            module Baz
+              extend T::Helpers
+
+              mixes_in_class_methods Foo::Bar
+              requires_ancestor { Foo::Bar }
+            end
           RB
 
-          assert_equal(contents, sorbet_sigs_to_rbs_comments(contents))
+          assert_equal(<<~RB, sorbet_sigs_to_rbs_comments(contents))
+            module Foo
+              extend T::Helpers
+
+              mixes_in_class_methods Bar
+
+              module Bar
+                def bar
+                end
+              end
+            end
+
+            # @requires_ancestor: Foo::Bar
+            module Baz
+              extend T::Helpers
+
+              mixes_in_class_methods Foo::Bar
+            end
+          RB
         end
 
         def test_translate_to_rbs_generics