Only translate top level assertions

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,23 +38,46 @@ 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
+          @rewriter << if node.name == :bind
+            Source::Replace.new(start_offset, end_offset - 1, rbs_annotation)
+          else
+            Source::Replace.new(start_offset, end_offset - 1, "#{dedent_value(node, value)} #{rbs_annotation}")
+          end
 
-        #: (Prism::CallNode) -> void
+          true
+        end
+
+        #: (Prism::CallNode) -> String
         def build_rbs_annotation(call)
           case call.name
           when :let
@@ -56,6 +88,10 @@ def build_rbs_annotation(call)
             srb_type = call.arguments&.arguments&.last #: as !nil
             rbs_type = RBI::Type.parse_node(srb_type).rbs_string
             "#: as #{rbs_type}"
+          when :bind
+            srb_type = call.arguments&.arguments&.last #: as !nil
+            rbs_type = RBI::Type.parse_node(srb_type).rbs_string
+            "#: self as #{rbs_type}"
           when :must
             "#: as !nil"
           when :unsafe
@@ -86,6 +122,8 @@ def t_annotation?(node)
           case node.name
           when :let, :cast
             return node.arguments&.arguments&.size == 2
+          when :bind
+            return node.arguments&.arguments&.size == 2 && node.arguments&.arguments&.first.is_a?(Prism::SelfNode)
           when :must, :unsafe
             return node.arguments&.arguments&.size == 1
           end

rbi/spoom.rbi

@@ -2934,23 +2934,26 @@ class Spoom::Sorbet::Translate::RBSCommentsToSorbetSigs < ::Spoom::Sorbet::Trans
 end
 
 class Spoom::Sorbet::Translate::SorbetAssertionsToRBSComments < ::Spoom::Sorbet::Translate::Translator
-  sig { override.params(node: ::Prism::CallNode).void }
-  def visit_call_node(node); end
-
   sig { override.params(node: ::Prism::IfNode).void }
   def visit_if_node(node); end
 
+  sig { override.params(node: ::Prism::StatementsNode).void }
+  def visit_statements_node(node); end
+
   private
 
   sig { params(node: ::Prism::Node).returns(T::Boolean) }
   def at_end_of_line?(node); end
 
-  sig { params(call: ::Prism::CallNode).void }
+  sig { params(call: ::Prism::CallNode).returns(::String) }
   def build_rbs_annotation(call); end
 
   sig { params(assign: ::Prism::Node, value: ::Prism::Node).returns(::String) }
   def dedent_value(assign, value); end
 
+  sig { params(node: ::Prism::Node).returns(T::Boolean) }
+  def maybe_translate_assertion(node); end
+
   sig { params(node: T.nilable(::Prism::Node)).returns(T::Boolean) }
   def t?(node); end
 

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)
@@ -490,6 +498,24 @@ def test_doesnt_translate_in_ternary_expressions
           assert_equal(rb, rbi_to_rbs(rb))
         end
 
+        def test_translate_bind
+          rb = <<~RB
+            T.bind(self, T.class_of(String))
+
+            T.bind(foo, String)
+
+            before_enqueue { T.bind(self, String) }
+          RB
+
+          assert_equal(<<~RB, rbi_to_rbs(rb))
+            #: self as singleton(String)
+
+            T.bind(foo, String)
+
+            before_enqueue { T.bind(self, String) }
+          RB
+        end
+
         private
 
         #: (String) -> String