Rewrite type aliases

Author: paracycle

lib/spoom/parse.rb

@@ -24,5 +24,22 @@ def parse_ruby(ruby, file:, comments: false)
 
       result.value
     end
+
+    #: (String ruby, file: String) -> [Prism::Node, Array[Prism::Comment]]
+    def parse_ruby_with_comments(ruby, file:)
+      result = Prism.parse(ruby)
+
+      unless result.success?
+        message = +"Error while parsing #{file}:\n"
+        result.errors.each do |e|
+          message << "- #{e.message} (at #{e.location.start_line}:#{e.location.start_column})\n"
+        end
+        raise ParseError, message
+      end
+
+      result.attach_comments!
+
+      [result.value, result.comments]
+    end
   end
 end

lib/spoom/rbs.rb

@@ -70,6 +70,7 @@ def initialize(string, location)
 
     class Annotation < Comment; end
     class Signature < Comment; end
+    class TypeAlias < Comment; end
 
     module ExtractRBSComments
       #: (Prism::Node) -> Comments

lib/spoom/sorbet/translate/rbs_comments_to_sorbet_sigs.rb

@@ -14,6 +14,16 @@ def initialize(ruby_contents, file:, max_line_length: nil)
           @max_line_length = max_line_length
         end
 
+        # @override
+        #: (Prism::ProgramNode node) -> void
+        def visit_program_node(node)
+          # Process all type aliases from the entire file first
+          apply_type_aliases(@comments)
+
+          # Now process the rest of the file with type aliases available
+          super
+        end
+
         # @override
         #: (Prism::ClassNode node) -> void
         def visit_class_node(node)
@@ -274,6 +284,77 @@ def already_extends?(node, constant_regex)
             true
           end
         end
+
+        #: (Array[Prism::Comment]) -> Array[RBS::TypeAlias]
+        def collect_type_aliases(comments)
+          type_aliases = [] #: Array[RBS::TypeAlias]
+
+          return type_aliases if comments.empty?
+
+          continuation_comments = [] #: Array[Prism::Comment]
+
+          comments.reverse_each do |comment|
+            string = comment.slice
+
+            if string.start_with?("#:")
+              string = string.delete_prefix("#:").strip
+              location = comment.location
+
+              if string.start_with?("type ")
+                continuation_comments.reverse_each do |continuation_comment|
+                  string = "#{string}#{continuation_comment.slice.delete_prefix("#|")}"
+                  location = location.join(continuation_comment.location)
+                end
+
+                type_aliases << Spoom::RBS::TypeAlias.new(string, location)
+              end
+
+              # Clear the continuation comments regardless of whether we found a type alias or not
+              continuation_comments.clear
+            elsif string.start_with?("#|")
+              continuation_comments << comment
+            else
+              continuation_comments.clear
+            end
+          end
+
+          type_aliases
+        end
+
+        #: (Array[Prism::Comment]) -> void
+        def apply_type_aliases(comments)
+          type_aliases = collect_type_aliases(comments)
+
+          type_aliases.each do |type_alias|
+            indent = " " * type_alias.location.start_column
+            insert_pos = adjust_to_line_start(type_alias.location.start_offset)
+
+            from = insert_pos
+            to = adjust_to_line_end(type_alias.location.end_offset)
+
+            *, decls = ::RBS::Parser.parse_signature(type_alias.string)
+
+            # We only expect there to be a single type alias declaration
+            next unless decls.size == 1 && decls.first.is_a?(::RBS::AST::Declarations::TypeAlias)
+
+            rbs_type = decls.first
+            sorbet_type = RBI::RBS::TypeTranslator.translate(rbs_type.type)
+
+            alias_name = ::RBS::TypeName.new(
+              namespace: rbs_type.name.namespace,
+              name: rbs_type.name.name.to_s.gsub(/(?:^|_)([a-z\d]*)/i) do |match|
+                match = match.delete_prefix("_")
+                !match.empty? ? match[0].upcase.concat(match[1..-1]) : +""
+              end,
+            )
+
+            @rewriter << Source::Delete.new(from, to)
+            @rewriter << Source::Insert.new(insert_pos, "#{indent}#{alias_name} = T.type_alias { #{sorbet_type.to_rbi} }\n")
+          rescue ::RBS::ParsingError, ::RBI::Error
+            # Ignore type aliases with errors
+            next
+          end
+        end
       end
     end
   end

lib/spoom/sorbet/translate/translator.rb

@@ -19,8 +19,9 @@ def initialize(ruby_contents, file:)
             ruby_contents.encode("UTF-8")
           end #: String
 
-          node = Spoom.parse_ruby(ruby_contents, file: file, comments: true)
+          node, comments = Spoom.parse_ruby_with_comments(ruby_contents, file: file)
           @node = node #: Prism::Node
+          @comments = comments #: Array[Prism::Comment]
           @ruby_bytes = ruby_contents.bytes #: Array[Integer]
           @rewriter = Spoom::Source::Rewriter.new #: Source::Rewriter
         end

rbi/spoom.rbi

@@ -2584,6 +2584,7 @@ module Spoom::RBS::ExtractRBSComments
 end
 
 class Spoom::RBS::Signature < ::Spoom::RBS::Comment; end
+class Spoom::RBS::TypeAlias < ::Spoom::RBS::Comment; end
 Spoom::SPOOM_PATH = T.let(T.unsafe(nil), String)
 module Spoom::Sorbet; end
 Spoom::Sorbet::BIN_PATH = T.let(T.unsafe(nil), String)

test/spoom/sorbet/translate/rbs_comments_to_sorbet_sigs_test.rb

@@ -466,6 +466,160 @@ def bar; end
           assert_equal(contents, rbs_comments_to_sorbet_sigs(contents))
         end
 
+        def test_translate_type_alias
+          contents = <<~RB
+            module Aliases
+              #: type foo = Integer | String
+              #: type multiLine =
+              #|   Integer |
+              #|   String
+            end
+
+            #: (Aliases::foo a) -> Aliases::multiLine
+            def bar(a)
+              42
+            end
+          RB
+
+          assert_equal(<<~RB, rbs_comments_to_sorbet_sigs(contents))
+            module Aliases
+              Foo = T.type_alias { T.any(Integer, String) }
+              MultiLine = T.type_alias { T.any(Integer, String) }
+            end
+
+            sig { params(a: Aliases::Foo).returns(Aliases::MultiLine) }
+            def bar(a)
+              42
+            end
+          RB
+        end
+
+        def test_translate_type_alias_with_complex_type
+          contents = <<~RB
+            #: type Foo::user_id = Integer
+            #: type ::Bar::user_data = { id: Foo::user_id, name: String }
+
+            #: (::Bar::user_data data) -> Foo::user_id
+            def process_user(data)
+              data[:id]
+            end
+          RB
+
+          assert_equal(<<~RB, rbs_comments_to_sorbet_sigs(contents))
+            Foo::UserId = T.type_alias { Integer }
+            ::Bar::UserData = T.type_alias { { id: Foo::UserId, name: String } }
+
+            sig { params(data: ::Bar::UserData).returns(Foo::UserId) }
+            def process_user(data)
+              data[:id]
+            end
+          RB
+        end
+
+        def test_translate_type_alias_in_class
+          contents = <<~RB
+            class Example
+              #: type status = :pending | :completed | :failed
+
+              #: () -> status
+              def get_status
+                :pending
+              end
+            end
+          RB
+
+          assert_equal(<<~RB, rbs_comments_to_sorbet_sigs(contents))
+            class Example
+              Status = T.type_alias { T.untyped }
+
+              sig { returns(Status) }
+              def get_status
+                :pending
+              end
+            end
+          RB
+        end
+
+        def test_translate_type_alias_with_generics
+          contents = <<~RB
+            #: type list = Array[Integer]
+
+            #: (list items) -> list
+            def double_items(items)
+              items.map { |x| x * 2 }
+            end
+          RB
+
+          assert_equal(<<~RB, rbs_comments_to_sorbet_sigs(contents))
+            List = T.type_alias { T::Array[Integer] }
+
+            sig { params(items: List).returns(List) }
+            def double_items(items)
+              items.map { |x| x * 2 }
+            end
+          RB
+        end
+
+        def test_translate_type_alias_with_union
+          contents = <<~RB
+            #: type nullable_string = String?
+
+            #: (nullable_string text) -> String
+            def ensure_string(text)
+              text || ""
+            end
+          RB
+
+          assert_equal(<<~RB, rbs_comments_to_sorbet_sigs(contents))
+            NullableString = T.type_alias { T.nilable(String) }
+
+            sig { params(text: NullableString).returns(String) }
+            def ensure_string(text)
+              text || ""
+            end
+          RB
+        end
+
+        def test_translate_type_alias_that_does_not_exist
+          contents = <<~RB
+            #: () -> notFound
+            def foo
+            end
+          RB
+
+          assert_equal(<<~RB, rbs_comments_to_sorbet_sigs(contents))
+            sig { returns(NotFound) }
+            def foo
+            end
+          RB
+        end
+
+        def test_translate_broken_type_alias_continuation
+          contents = <<~RB
+            #: type multiLine =
+            #| String
+            #| | Integer
+            # foo bar baz
+            #| | Symbol
+
+            #: () -> multiLine
+            def foo
+              ""
+            end
+          RB
+
+          assert_equal(<<~RB, rbs_comments_to_sorbet_sigs(contents))
+            MultiLine = T.type_alias { T.any(String, Integer) }
+            # foo bar baz
+            #| | Symbol
+
+            sig { returns(MultiLine) }
+            def foo
+              ""
+            end
+          RB
+        end
+
         private
 
         #: (String, ?max_line_length: Integer?) -> String