Add gem extensions

README.md

@@ -41,6 +41,7 @@ Tapioca makes it easy to work with [Sorbet](https://sorbet.org) in your codebase
     * [Changing the strictness level of the RBI for a gem](#changing-the-strictness-level-of-the-rbi-for-a-gem)
     * [Keeping RBI files for gems up-to-date](#keeping-rbi-files-for-gems-up-to-date)
     * [Importing hand written signatures from gem's `rbi/` folder](#importing-hand-written-signatures-from-gems-rbi-folder)
+    * [Writing custom gem extensions](#writing-custom-gem-extensions)
   * [Pulling RBI annotations from remote sources](#pulling-rbi-annotations-from-remote-sources)
     * [Basic authentication](#basic-authentication)
     * [Using a .netrc file](#using-a-netrc-file)
@@ -345,6 +346,55 @@ This option can be used in CI to make sure the RBI files are *up-to-date* and en
 
 Tapioca will import any signatures found in the `rbi/` folder of a given gem and combine them with the RBIs it generates. This is useful when a gem doesn't want to depend on `sorbet-runtime` but still wants to provide type safety to users during static checks. Note that the `rbi/` folder needs to be included in the gem release using the `.gemspec` file. Applications can choose not to import these signatures using the `--no-exported-gem-rbis` flag.
 
+#### Writing custom gem extensions
+
+Similar to [DSL extensions](#writing-custom-dsl-extensions), it is sometimes useful to patch and modify certain behavior
+from gems to ensure that RBIs are generated correctly. Gem RBIs are generated by reflecting over what the gem defines at
+runtime, which means that anything the gem never really defines is invisible to Tapioca. A common example is methods
+handled by `method_missing`:
+
+```ruby
+# my_gem/lib/my_gem/settings.rb
+
+module MyGem
+  class Settings
+    SETTINGS = [:host, :port, :timeout]
+
+    def method_missing(name, *args)
+      SETTINGS.include?(name) ? @config[name] : super
+    end
+
+    def respond_to_missing?(name, include_private = false)
+      SETTINGS.include?(name) || super
+    end
+  end
+end
+```
+
+Since `host`, `port` and `timeout` are never defined as real methods, they will be missing from the generated RBI. An
+extension can define them, so that Tapioca is able to see them:
+
+```ruby
+# my_gem/lib/tapioca/gem/extensions/my_gem.rb
+
+require "my_gem"
+
+module MyGem
+  class Settings
+    SETTINGS.each do |name|
+      define_method(name) { @config[name] }
+    end
+  end
+end
+```
+
+Extensions are only loaded during RBI generation, so they never change the behavior of the gem for applications using
+it. They are also loaded before the bundle is required, which is why the extension has to require what it wants to
+patch.
+
+Applications can also define their own gem extensions by placing them inside the `sorbet/tapioca/gem/extensions`
+directory.
+
 ### Pulling RBI annotations from remote sources
 
 Since Tapioca does not perform any type inference, the RBI files generated for the gems do not contain any type signatures. Instead, Tapioca relies on the community to provide high-quality, manually written RBI annotations for public gems.

lib/tapioca/gemfile.rb

@@ -48,8 +48,21 @@ def require_bundle
       end
     end
 
+    #: (String path) -> bool
+    def excluded_gem_path?(path)
+      excluded_gem_specs.any? { |spec| spec.contains_path?(path) }
+    end
+
     private
 
+    #: -> Array[GemSpec]
+    def excluded_gem_specs
+      @excluded_gem_specs ||= @excluded_gems.filter_map do |name|
+        spec = ::Gem.loaded_specs[name]
+        GemSpec.new(spec) if spec
+      end #: Array[GemSpec]?
+    end
+
     #: File
     attr_reader(:gemfile, :lockfile)
 

lib/tapioca/loaders/gem.rb

@@ -49,6 +49,39 @@ def initialize(bundle:, prerequire:, postrequire:, default_command:, halt_upon_l
         @halt_upon_load_error = halt_upon_load_error
       end
 
+      #: -> void
+      def load_gem_extensions
+        say("Loading gem extension classes... ")
+
+        # Extensions are loaded before the bundle is required so that they can patch the gems
+        # they apply to as those gems are being loaded.
+        Dir.glob("#{Tapioca::TAPIOCA_DIR}/gem/extensions/**/*.rb").each do |extension|
+          require File.expand_path(extension)
+        end
+
+        ::Gem.find_files("tapioca/gem/extensions/*.rb").each do |extension|
+          next if @bundle.excluded_gem_path?(extension)
+
+          require File.expand_path(extension)
+        end
+
+        say("Done", :green)
+      end
+
+      #: (Tapioca::Gemfile gemfile, String? initialize_file, String? require_file, bool halt_upon_load_error) -> void
+      def load_bundle(gemfile, initialize_file, require_file, halt_upon_load_error)
+        require_helper(initialize_file)
+        load_gem_extensions
+
+        load_rails_application(halt_upon_load_error: halt_upon_load_error)
+
+        gemfile.require_bundle
+
+        require_helper(require_file)
+
+        load_rails_engines
+      end
+
       #: -> void
       def require_gem_file
         say("Requiring all gems to prepare for compiling... ")

lib/tapioca/loaders/loader.rb

@@ -15,19 +15,6 @@ def load = raise NotImplementedError, "Abstract method called"
 
       private
 
-      #: (Tapioca::Gemfile gemfile, String? initialize_file, String? require_file, bool halt_upon_load_error) -> void
-      def load_bundle(gemfile, initialize_file, require_file, halt_upon_load_error)
-        require_helper(initialize_file)
-
-        load_rails_application(halt_upon_load_error: halt_upon_load_error)
-
-        gemfile.require_bundle
-
-        require_helper(require_file)
-
-        load_rails_engines
-      end
-
       #: (?environment_load: bool, ?eager_load: bool, ?app_root: String, ?halt_upon_load_error: bool) -> void
       def load_rails_application(environment_load: false, eager_load: false, app_root: ".", halt_upon_load_error: true)
         return unless File.exist?(File.expand_path("config/application.rb", app_root))

spec/tapioca/cli/gem_spec.rb

@@ -2298,6 +2298,142 @@ class Application < Rails::Application
           assert_success_status(result)
         end
       end
+
+      describe "custom extensions" do
+        after do
+          project.write_gemfile!(project.tapioca_gemfile)
+          @project.require_default_gems
+          @project.remove!("sorbet/rbi")
+          @project.remove!("sorbet/tapioca/gem")
+          @project.remove!("../gems")
+        end
+
+        it "loads extensions" do
+          foo = mock_gem("foo", "0.0.1") do
+            write!("lib/foo.rb", <<~RUBY)
+              module Patch
+                def [](*types)
+                  self
+                end
+              end
+
+              class Foo
+                extend T::Generic
+                Value = type_member
+                extend Patch
+
+                sig do
+                  type_parameters(:Value).
+                    params(
+                      block: T.proc.returns(T.type_parameter(:Value))
+                    ).returns(Foo[T.type_parameter(:Value)])
+                end
+                def something(&block); end
+              end
+            RUBY
+
+            write!("lib/tapioca/gem/extensions/foo.rb", <<~RUBY)
+              require "foo"
+
+              module Patch
+                def [](*types)
+                  super
+                end
+              end
+            RUBY
+          end
+
+          @project.require_mock_gem(foo)
+          @project.bundle_install!
+
+          result = @project.tapioca("gem foo")
+
+          assert_stdout_includes(result, "Loading gem extension classes... Done")
+
+          assert_project_file_includes("sorbet/rbi/gems/foo@0.0.1.rbi", <<~RBI)
+            class Foo
+              extend T::Generic
+              extend ::Patch
+
+              Value = type_member
+
+              sig do
+                type_parameters(:Value)
+                  .params(
+                    block: T.proc.returns(T.type_parameter(:Value))
+                  ).returns(Foo[T.type_parameter(:Value)])
+              end
+              def something(&block); end
+            end
+
+            module Patch
+              def [](*types); end
+            end
+          RBI
+
+          assert_empty_stderr(result)
+          assert_success_status(result)
+        end
+
+        it "loads extensions defined by the project" do
+          foo = mock_gem("foo", "0.0.1") do
+            write!("lib/foo.rb", <<~RUBY)
+              class Foo
+                SETTINGS.each { |setting| define_method(setting) { nil } } if defined?(SETTINGS)
+              end
+            RUBY
+          end
+
+          @project.require_mock_gem(foo)
+          @project.bundle_install!
+
+          @project.write!("sorbet/tapioca/gem/extensions/foo.rb", <<~RUBY)
+            SETTINGS = [:host, :port]
+          RUBY
+
+          result = @project.tapioca("gem foo")
+
+          assert_project_file_includes("sorbet/rbi/gems/foo@0.0.1.rbi", <<~RBI)
+            class Foo
+              def host; end
+              def port; end
+            end
+          RBI
+
+          assert_empty_stderr(result)
+          assert_success_status(result)
+        end
+
+        it "does not load extensions from excluded gems" do
+          foo = mock_gem("foo", "0.0.1") do
+            write!("lib/foo.rb", "module Foo; end")
+
+            write!("lib/tapioca/gem/extensions/foo.rb", <<~RUBY)
+              puts "FOO EXTENSION LOADED"
+            RUBY
+          end
+
+          bar = mock_gem("bar", "0.0.1") do
+            write!("lib/bar.rb", "module Bar; end")
+
+            write!("lib/tapioca/gem/extensions/bar.rb", <<~RUBY)
+              puts "BAR EXTENSION LOADED"
+            RUBY
+          end
+
+          @project.require_mock_gem(foo)
+          @project.require_mock_gem(bar)
+          @project.bundle_install!
+
+          result = @project.tapioca("gem bar --exclude foo")
+
+          assert_stdout_includes(result, "BAR EXTENSION LOADED")
+          refute_includes(result.out, "FOO EXTENSION LOADED", result.to_s)
+
+          assert_empty_stderr(result)
+          assert_success_status(result)
+        end
+      end
     end
   end
 end