Only translate top level assertions

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

Author: Morriar

lib/spoom/sorbet/translate/sorbet_assertions_to_rbs_comments.rb

@@ -8,6 +8,15 @@ module Translate
       class SorbetAssertionsToRBSComments < Translator
         LINE_BREAK = "\n".ord #: Integer
 
+        # @override
+        #: (Prism::StatementsNode) -> void
+        def visit_statements_node(node)
+          node.body.each do |statement|
+            translated = maybe_translate_assertion(statement)
+            visit(statement) unless translated
+          end
+        end
+
         # @override
         #: (Prism::IfNode) -> void
         def visit_if_node(node)
@@ -29,21 +38,39 @@ def visit_if_node(node)
           end
         end
 
-        # @override
-        #: (Prism::CallNode) -> void
-        def visit_call_node(node)
-          return super unless t_annotation?(node)
-          return super unless at_end_of_line?(node)
+        private
+
+        #: (Prism::Node) -> bool
+        def maybe_translate_assertion(node)
+          node = case node
+          when Prism::MultiWriteNode,
+               Prism::ClassVariableWriteNode, Prism::ClassVariableAndWriteNode, Prism::ClassVariableOperatorWriteNode, Prism::ClassVariableOrWriteNode,
+               Prism::ConstantWriteNode, Prism::ConstantAndWriteNode, Prism::ConstantOperatorWriteNode, Prism::ConstantOrWriteNode,
+               Prism::ConstantPathWriteNode, Prism::ConstantPathAndWriteNode, Prism::ConstantPathOperatorWriteNode, Prism::ConstantPathOrWriteNode,
+               Prism::GlobalVariableWriteNode, Prism::GlobalVariableAndWriteNode, Prism::GlobalVariableOperatorWriteNode, Prism::GlobalVariableOrWriteNode,
+               Prism::InstanceVariableWriteNode, Prism::InstanceVariableAndWriteNode, Prism::InstanceVariableOperatorWriteNode, Prism::InstanceVariableOrWriteNode,
+               Prism::LocalVariableWriteNode, Prism::LocalVariableAndWriteNode, Prism::LocalVariableOperatorWriteNode, Prism::LocalVariableOrWriteNode,
+               Prism::CallAndWriteNode, Prism::CallOperatorWriteNode, Prism::CallOrWriteNode
+            node.value
+          when Prism::CallNode
+            node
+          else
+            return false
+          end
+
+          return false unless node.is_a?(Prism::CallNode)
+          return false unless t_annotation?(node)
+          return false unless at_end_of_line?(node)
 
           value = T.must(node.arguments&.arguments&.first)
           rbs_annotation = build_rbs_annotation(node)
 
           start_offset = node.location.start_offset
           end_offset = node.location.end_offset
           @rewriter << Source::Replace.new(start_offset, end_offset - 1, "#{dedent_value(node, value)} #{rbs_annotation}")
-        end
 
-        private
+          true
+        end
 
         #: (Prism::CallNode) -> void
         def build_rbs_annotation(call)

test/spoom/sorbet/translate/sorbet_assertions_to_rbs_comments_test.rb

@@ -482,6 +482,14 @@ def test_doesnt_translate_cast_in_string_interpolation
           assert_equal(rb, rbi_to_rbs(rb))
         end
 
+        def test_doesnt_translate_in_expressions
+          rb = <<~RB
+            a - T.must(b)
+          RB
+
+          assert_equal(rb, rbi_to_rbs(rb))
+        end
+
         def test_doesnt_translate_in_ternary_expressions
           rb = <<~RB
             a = T.must(b) ? T.must(c) : T.must(d)