diff --git a/lib/spoom/sorbet/translate/sorbet_assertions_to_rbs_comments.rb b/lib/spoom/sorbet/translate/sorbet_assertions_to_rbs_comments.rb index 84856a81..406bbdff 100644 --- a/lib/spoom/sorbet/translate/sorbet_assertions_to_rbs_comments.rb +++ b/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 diff --git a/rbi/spoom.rbi b/rbi/spoom.rbi index 2777aef6..6ccce9b1 100644 --- a/rbi/spoom.rbi +++ b/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 diff --git a/test/spoom/sorbet/translate/sorbet_assertions_to_rbs_comments_test.rb b/test/spoom/sorbet/translate/sorbet_assertions_to_rbs_comments_test.rb index 43bdc8a0..6ec81251 100644 --- a/test/spoom/sorbet/translate/sorbet_assertions_to_rbs_comments_test.rb +++ b/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