Add support for gem extensions

README.md

@@ -24,6 +24,7 @@
   - [Setup](#setup)
     - [Generic](#generic)
     - [for_gem](#for_gem)
+    - [for_gem_extension](#for_gem_extension)
   - [Autoloading](#autoloading)
   - [Eager loading](#eager-loading)
     - [Eager load exclusions](#eager-load-exclusions)
@@ -410,6 +411,42 @@ Otherwise, there's a flag to say the extra stuff is OK:
 Zeitwerk::Loader.for_gem(warn_on_extra_files: false)
 ```
 
+<a id="markdown-for_gem_extension" name="for_gem_extension"></a>
+#### for_gem_extension
+
+Let's suppose you are writing a gem to extend `Net::HTTP` with some niche feature. By [convention](https://guides.rubygems.org/name-your-gem/):
+
+* The gem should be called `net-http-niche_feature`. That is, dashes for the extended part, a dash, and underscores for yours.
+* The namespace should be `Net::HTTP::NicheFeature`.
+* The entry point should be `lib/net/http/niche_feature.rb`.
+* Optionally, the gem could have a top-level `lib/net-http-niche_feature.rb`, but, if defined, that one should have just a `require` call for the entry point.
+
+Gem extensions following the conventions above have a dedicated loader constructor:
+
+```ruby
+# lib/net/http/niche_feature.rb
+
+require "net/http"
+require "zeitwerk"
+
+loader = Zeitwerk::Loader.for_gem_extension(Net::HTTP)
+loader.setup
+
+module Net::HTTP::NicheFeature
+  # Since the setup has been performed, at this point we are already able
+  # to reference project constants, in this case Net::HTTP::NicheFeature::MyMixin.
+  include MyMixin
+end
+```
+
+`for_gem_extension` expects as argument the the namespace being extended, which has to be a class or module object.
+
+The file `lib/net/http/niche_feature/version.rb` is expected to define `Net::HTTP::NicheFeature::VERSION`.
+
+Due to technical reasons, the entry point of the gem has to be loaded with `Kernel#require`. Loading that file with `Kernel#load` or `Kernel#require_relative` won't generally work. This is important if you load the entry point from the optional dasherized top-level file.
+
+`Zeitwerk::Loader.for_gem_extension` is idempotent when invoked from the same file, to support gems that want to reload (unlikely).
+
 <a id="markdown-autoloading" name="autoloading"></a>
 ### Autoloading
 

lib/zeitwerk/gem_inflector.rb

@@ -5,8 +5,8 @@ class GemInflector < Inflector
     # @sig (String) -> void
     def initialize(root_file)
       namespace     = File.basename(root_file, ".rb")
-      lib_dir       = File.dirname(root_file)
-      @version_file = File.join(lib_dir, namespace, "version.rb")
+      root_dir      = File.dirname(root_file)
+      @version_file = File.join(root_dir, namespace, "version.rb")
     end
 
     # @sig (String, String) -> String

lib/zeitwerk/gem_loader.rb

@@ -3,27 +3,31 @@
 module Zeitwerk
   # @private
   class GemLoader < Loader
+    include RealModName
+
     # Users should not create instances directly, the public interface is
     # `Zeitwerk::Loader.for_gem`.
     private_class_method :new
 
     # @private
     # @sig (String, bool) -> Zeitwerk::GemLoader
-    def self._new(root_file, warn_on_extra_files:)
-      new(root_file, warn_on_extra_files: warn_on_extra_files)
+    def self._new(root_file, namespace:, warn_on_extra_files:)
+      new(root_file, namespace: namespace, warn_on_extra_files: warn_on_extra_files)
     end
 
     # @sig (String, bool) -> void
-    def initialize(root_file, warn_on_extra_files:)
+    def initialize(root_file, namespace:, warn_on_extra_files:)
       super()
 
-      @tag                 = File.basename(root_file, ".rb")
+      @tag = File.basename(root_file, ".rb")
+      @tag = real_mod_name(namespace) + "-" + @tag unless namespace.equal?(Object)
+
       @inflector           = GemInflector.new(root_file)
       @root_file           = File.expand_path(root_file)
-      @lib                 = File.dirname(root_file)
+      @root_dir            = File.dirname(root_file)
       @warn_on_extra_files = warn_on_extra_files
 
-      push_dir(@lib)
+      push_dir(@root_dir, namespace: namespace)
     end
 
     # @sig () -> void
@@ -38,7 +42,7 @@ def setup
     def warn_on_extra_files
       expected_namespace_dir = @root_file.delete_suffix(".rb")
 
-      ls(@lib) do |basename, abspath|
+      ls(@root_dir) do |basename, abspath|
         next if abspath == @root_file
         next if abspath == expected_namespace_dir
 

lib/zeitwerk/loader.rb

@@ -270,6 +270,7 @@ class << self
       # This is a shortcut for
       #
       #   require "zeitwerk"
+      #
       #   loader = Zeitwerk::Loader.new
       #   loader.tag = File.basename(__FILE__, ".rb")
       #   loader.inflector = Zeitwerk::GemInflector.new(__FILE__)
@@ -284,7 +285,28 @@ class << self
       # @sig (bool) -> Zeitwerk::GemLoader
       def for_gem(warn_on_extra_files: true)
         called_from = caller_locations(1, 1).first.path
-        Registry.loader_for_gem(called_from, warn_on_extra_files: warn_on_extra_files)
+        Registry.loader_for_gem(called_from, namespace: Object, warn_on_extra_files: warn_on_extra_files)
+      end
+
+      # This is a shortcut for
+      #
+      #   require "zeitwerk"
+      #
+      #   loader = Zeitwerk::Loader.new
+      #   loader.tag = namespace.name + "-" + File.basename(__FILE__, ".rb")
+      #   loader.inflector = Zeitwerk::GemInflector.new(__FILE__)
+      #   loader.push_dir(__dir__, namespace: namespace)
+      #
+      # except that this method returns the same object in subsequent calls from
+      # the same file, in the unlikely case the gem wants to be able to reload.
+      #
+      # This method returns a subclass of Zeitwerk::Loader, but the exact type
+      # is private, client code can only rely on the interface.
+      #
+      # @sig (bool) -> Zeitwerk::GemLoader
+      def for_gem_extension(namespace)
+        called_from = caller_locations(1, 1).first.path
+        Registry.loader_for_gem(called_from, namespace: namespace, warn_on_extra_files: false)
       end
 
       # Broadcasts `eager_load` to all loaders. Those that have not been setup

lib/zeitwerk/registry.rb

@@ -86,8 +86,8 @@ def unregister_loader(loader)
       #
       # @private
       # @sig (String) -> Zeitwerk::Loader
-      def loader_for_gem(root_file, warn_on_extra_files:)
-        gem_loaders_by_root_file[root_file] ||= GemLoader._new(root_file, warn_on_extra_files: warn_on_extra_files)
+      def loader_for_gem(root_file, namespace:, warn_on_extra_files:)
+        gem_loaders_by_root_file[root_file] ||= GemLoader._new(root_file, namespace: namespace, warn_on_extra_files: warn_on_extra_files)
       end
 
       # @private

test/lib/zeitwerk/test_for_gem_extension.rb

@@ -0,0 +1,94 @@
+# frozen_string_literal: true
+
+require "test_helper"
+
+class TestForGemExtension < LoaderTest
+  ROOT_DIR = "lib/test_for_gem_extension"
+  MY_GEM_EXTENSION = ["#{ROOT_DIR}/my_gem_extension.rb", <<~EOS]
+    $for_gem_extension_test_loader = Zeitwerk::Loader.for_gem_extension(#{self})
+    $for_gem_extension_test_loader.enable_reloading
+    $for_gem_extension_test_loader.setup
+
+    module #{self}::MyGemExtension
+    end
+  EOS
+
+  def with_my_gem_extension(files = [MY_GEM_EXTENSION], rq = true)
+    with_files(files) do
+      with_load_path("lib") do
+        if rq
+          assert require("test_for_gem_extension/my_gem_extension")
+          assert self.class::MyGemExtension
+        end
+        yield
+      end
+    end
+  end
+
+  test "sets things correctly" do
+    files = [
+      MY_GEM_EXTENSION,
+      ["#{ROOT_DIR}/my_gem_extension/foo.rb", "class #{self.class}::MyGemExtension::Foo; end"],
+      ["#{ROOT_DIR}/my_gem_extension/foo/bar.rb", "#{self.class}::MyGemExtension::Foo::Bar = true"]
+    ]
+    with_my_gem_extension(files) do
+      assert self.class::MyGemExtension::Foo::Bar
+
+      $for_gem_extension_test_loader.unload
+      assert !self.class.const_defined?(:MyGemExtension)
+
+      $for_gem_extension_test_loader.setup
+      assert self.class::MyGemExtension::Foo::Bar
+    end
+  end
+
+  test "is idempotent" do
+    $for_gem_extension_zs = []
+    files = [
+      ["#{ROOT_DIR}/my_gem_extension.rb", <<-EOS],
+        $for_gem_extension_zs << Zeitwerk::Loader.for_gem_extension(#{self.class})
+        $for_gem_extension_zs.last.enable_reloading
+        $for_gem_extension_zs.last.setup
+
+        module #{self.class}::MyGemExtension
+        end
+      EOS
+    ]
+
+    with_my_gem_extension(files) do
+      $for_gem_extension_zs.first.unload
+      assert !self.class.const_defined?(:MyGemExtension)
+
+      $for_gem_extension_zs.first.setup
+      assert self.class::MyGemExtension
+
+      assert_equal 2, $for_gem_extension_zs.size
+      assert_same $for_gem_extension_zs.first, $for_gem_extension_zs.last
+    end
+  end
+
+  test "configures the gem inflector by default" do
+    files = [MY_GEM_EXTENSION, ["#{ROOT_DIR}/my_gem_extension/version.rb", "#{self.class}::MyGemExtension::VERSION = '1.0'"]]
+    with_my_gem_extension(files) do
+      assert_instance_of Zeitwerk::GemInflector, $for_gem_extension_test_loader.inflector
+      assert_equal "1.0", self.class::MyGemExtension::VERSION
+    end
+  end
+
+  test "configures the namespace plus basename of the root file as loader tag" do
+    with_my_gem_extension do
+      assert_equal "#{self.class}-my_gem_extension", $for_gem_extension_test_loader.tag
+    end
+  end
+
+  test "works too if going through a dasherized entry point (require)" do
+    files = [
+      ["lib/my-gem-extension.rb", "require 'test_for_gem_extension/my_gem_extension'"],
+      MY_GEM_EXTENSION,
+    ]
+    with_my_gem_extension(files, false) do
+      assert require("my-gem-extension")
+      assert self.class::MyGemExtension
+    end
+  end
+end