Merge pull request #791 from Shopify/uk-type-alias-rewriting

KaanOzkan · web-flow · commit 3f4ce874bdab · 2025-10-17T16:24:41.000-04:00
Rewrite type aliases
diff --git a/lib/spoom/deadcode/index.rb b/lib/spoom/deadcode/index.rb
@@ -49,7 +49,7 @@ def index_erb(erb, file:, plugins: [])
 
       #: (String rb, file: String, ?plugins: Array[Plugins::Base]) -> void
       def index_ruby(rb, file:, plugins: [])
-        node = Spoom.parse_ruby(rb, file: file, comments: true)
+        node, _ = Spoom.parse_ruby(rb, file: file)
 
         # Index definitions
         model_builder = Model::Builder.new(@model, file)
diff --git a/lib/spoom/parse.rb b/lib/spoom/parse.rb
@@ -7,22 +7,21 @@ module Spoom
   class ParseError < Error; end
 
   class << self
-    #: (String ruby, file: String, ?comments: bool) -> Prism::Node
-    def parse_ruby(ruby, file:, comments: false)
+    #: (String ruby, file: String) -> [Prism::Node, Array[Prism::Comment]]
+    def parse_ruby(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! if comments
+      result.attach_comments!
 
-      result.value
+      [result.value, result.comments]
     end
   end
 end
diff --git a/lib/spoom/rbs.rb b/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
diff --git a/lib/spoom/sorbet/metrics/code_metrics_visitor.rb b/lib/spoom/sorbet/metrics/code_metrics_visitor.rb
@@ -13,7 +13,7 @@ def collect_code_metrics(files)
             counters.increment("files")
 
             content = File.read(file)
-            node = Spoom.parse_ruby(content, file: file, comments: true)
+            node, _ = Spoom.parse_ruby(content, file: file)
             visitor = CodeMetricsVisitor.new(counters)
             visitor.visit(node)
           end
diff --git a/lib/spoom/sorbet/translate/rbs_comments_to_sorbet_sigs.rb b/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,78 @@ 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)
+            content = "#{indent}#{alias_name} = T.type_alias { #{sorbet_type.to_rbi} }\n"
+            @rewriter << Source::Insert.new(insert_pos, content)
+          rescue ::RBS::ParsingError, ::RBI::Error
+            # Ignore type aliases with errors
+            next
+          end
+        end
       end
     end
   end
diff --git a/lib/spoom/sorbet/translate/translator.rb b/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(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
diff --git a/rbi/spoom.rbi b/rbi/spoom.rbi
@@ -6,8 +6,8 @@
 
 module Spoom
   class << self
-    sig { params(ruby: ::String, file: ::String, comments: T::Boolean).returns(::Prism::Node) }
-    def parse_ruby(ruby, file:, comments: T.unsafe(nil)); end
+    sig { params(ruby: ::String, file: ::String).returns([::Prism::Node, T::Array[::Prism::Comment]]) }
+    def parse_ruby(ruby, file:); end
   end
 end
 
@@ -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)
@@ -2893,6 +2894,9 @@ class Spoom::Sorbet::Translate::RBSCommentsToSorbetSigs < ::Spoom::Sorbet::Trans
   sig { override.params(node: ::Prism::ModuleNode).void }
   def visit_module_node(node); end
 
+  sig { override.params(node: ::Prism::ProgramNode).void }
+  def visit_program_node(node); end
+
   sig { override.params(node: ::Prism::SingletonClassNode).void }
   def visit_singleton_class_node(node); end
 
@@ -2912,6 +2916,12 @@ class Spoom::Sorbet::Translate::RBSCommentsToSorbetSigs < ::Spoom::Sorbet::Trans
   sig { params(annotations: T::Array[::Spoom::RBS::Annotation], sig: ::RBI::Sig).void }
   def apply_member_annotations(annotations, sig); end
 
+  sig { params(comments: T::Array[::Prism::Comment]).void }
+  def apply_type_aliases(comments); end
+
+  sig { params(comments: T::Array[::Prism::Comment]).returns(T::Array[::Spoom::RBS::TypeAlias]) }
+  def collect_type_aliases(comments); end
+
   sig { params(def_node: ::Prism::DefNode, comments: ::Spoom::RBS::Comments).void }
   def rewrite_def(def_node, comments); end
 
diff --git a/test/spoom/model/builder_test.rb b/test/spoom/model/builder_test.rb
@@ -464,7 +464,7 @@ def m1; end
       #: (String rb) -> Model
       def model(rb)
         file = "foo.rb"
-        node = Spoom.parse_ruby(rb, file: file, comments: true)
+        node, _ = Spoom.parse_ruby(rb, file: file)
         model = Model.new
         builder = Builder.new(model, file)
         builder.visit(node)
diff --git a/test/spoom/model/model_test.rb b/test/spoom/model/model_test.rb
@@ -104,7 +104,7 @@ module E; end
 
       #: (String rb) -> Model
       def model(rb)
-        node = Spoom.parse_ruby(rb, file: "foo.rb", comments: true)
+        node, _ = Spoom.parse_ruby(rb, file: "foo.rb")
 
         model = Model.new
         builder = Builder.new(model, "foo.rb")
diff --git a/test/spoom/model/namespace_visitor_test.rb b/test/spoom/model/namespace_visitor_test.rb
@@ -101,7 +101,7 @@ module ::M9::M10; end
 
       #: (String rb) -> Hash[String, String]
       def namespaces_for_locs(rb)
-        node = Spoom.parse_ruby(rb, file: "foo.rb")
+        node, _ = Spoom.parse_ruby(rb, file: "foo.rb")
 
         visitor = NamespacesForLocs.new
         visitor.visit(node)
diff --git a/test/spoom/model/references_visitor_test.rb b/test/spoom/model/references_visitor_test.rb
@@ -296,7 +296,7 @@ def test_visit_method_operators
       private
 
       def visit(code)
-        node = Spoom.parse_ruby(code, file: "-")
+        node, _ = Spoom.parse_ruby(code, file: "-")
 
         v = ReferencesVisitor.new("-")
         v.visit(node)
diff --git a/test/spoom/sorbet/translate/rbs_comments_to_sorbet_sigs_test.rb b/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
