Merge pull request #352 from Shopify/at-top-level-rewriter

Introduce NestTopLevelMembers rewriter

lib/rbi.rb

@@ -23,6 +23,7 @@ class Error < StandardError
 require "rbi/rewriters/merge_trees"
 require "rbi/rewriters/nest_singleton_methods"
 require "rbi/rewriters/nest_non_public_methods"
+require "rbi/rewriters/nest_top_level_members"
 require "rbi/rewriters/group_nodes"
 require "rbi/rewriters/remove_known_definitions"
 require "rbi/rewriters/attr_to_methods"

lib/rbi/rewriters/nest_top_level_members.rb

@@ -0,0 +1,68 @@
+# typed: strict
+# frozen_string_literal: true
+
+module RBI
+  module Rewriters
+    # This rewriter moves top-level members into a top-level Object class
+    #
+    # Example:
+    # ~~~rb
+    # def foo; end
+    # attr_reader :bar
+    # ~~~
+    #
+    # will be rewritten to:
+    #
+    # ~~~rb
+    # class Object
+    #  def foo; end
+    #  attr_reader :bar
+    # end
+    # ~~~
+    class NestTopLevelMembers < Visitor
+      extend T::Sig
+
+      sig { void }
+      def initialize
+        super
+
+        @top_level_object_class = T.let(nil, T.nilable(Class))
+      end
+
+      sig { override.params(node: T.nilable(Node)).void }
+      def visit(node)
+        return unless node
+
+        case node
+        when Tree
+          visit_all(node.nodes.dup)
+        else
+          scope = node.parent_scope
+          unless scope
+            parent = node.parent_tree
+            raise unless parent
+
+            node.detach
+
+            unless @top_level_object_class
+              @top_level_object_class = Class.new("Object")
+              parent.nodes << @top_level_object_class
+            end
+
+            @top_level_object_class << node
+          end
+        end
+      end
+    end
+  end
+
+  class Tree
+    extend T::Sig
+
+    sig { void }
+    def nest_top_level_members!
+      visitor = Rewriters::NestTopLevelMembers.new
+      visitor.visit(self)
+    end
+  end
+end

test/rbi/rewriters/nest_top_level_members_test.rb

@@ -0,0 +1,89 @@
+# typed: true
+# frozen_string_literal: true
+
+require "test_helper"
+
+module RBI
+  class NestTopLevelMembersTest < Minitest::Test
+    include TestHelper
+
+    def test_nest_top_level_members
+      tree = parse_rbi(<<~RBI)
+        module Foo
+          def m1; end
+        end
+
+        def m2; end
+        def self.m3; end
+        attr_reader :foo
+        send!
+
+        class << self; end
+      RBI
+
+      tree.nest_top_level_members!
+
+      assert_equal(<<~RBI, tree.string)
+        module Foo
+          def m1; end
+        end
+
+        class << self; end
+
+        class Object
+          def m2; end
+          def self.m3; end
+          attr_reader :foo
+          send!
+        end
+      RBI
+    end
+
+    def test_nest_top_level_members_reuse_the_same_class
+      tree = parse_rbi(<<~RBI)
+        def foo; end
+        class Foo; end
+        def bar; end
+      RBI
+
+      tree.nest_top_level_members!
+
+      assert_equal(<<~RBI, tree.string)
+        class Foo; end
+
+        class Object
+          def foo; end
+          def bar; end
+        end
+      RBI
+    end
+
+    def test_nest_top_level_members_duplicate_existing_object_class
+      tree = parse_rbi(<<~RBI)
+        class Object
+          def foo; end
+        end
+        def bar; end
+        class Object
+          def baz; end
+        end
+      RBI
+
+      tree.nest_top_level_members!
+
+      assert_equal(<<~RBI, tree.string)
+        class Object
+          def foo; end
+        end
+
+        class Object
+          def baz; end
+        end
+
+        class Object
+          def bar; end
+        end
+      RBI
+    end
+  end
+end