Add Zeitwerk support (#85)

This PR adds support for using Zeitwerk which:
- supports using collapsed directories
- for Rails applications with Zeitwerk, should speed up annotation
because it no longer eager loads

Author: drwl

lib/annotate_rb/eager_loader.rb

@@ -8,8 +8,12 @@ def call(options)
         options[:require].count > 0 && options[:require].each { |path| require path }
 
         if defined?(::Rails::Application)
-          klass = ::Rails::Application.send(:subclasses).first
-          klass.eager_load!
+          if defined?(::Zeitwerk)
+            # Delegate to Zeitwerk to load stuff as needed
+          else
+            klass = ::Rails::Application.send(:subclasses).first
+            klass.eager_load!
+          end
         else
           options[:model_dir].each do |dir|
             ::Rake::FileList["#{dir}/**/*.rb"].each do |fname|

lib/annotate_rb/model_annotator.rb

@@ -25,5 +25,6 @@ module ModelAnnotator
     autoload :ProjectAnnotationRemover, "annotate_rb/model_annotator/project_annotation_remover"
     autoload :AnnotatedFile, "annotate_rb/model_annotator/annotated_file"
     autoload :FileParser, "annotate_rb/model_annotator/file_parser"
+    autoload :ZeitwerkClassGetter, "annotate_rb/model_annotator/zeitwerk_class_getter"
   end
 end

lib/annotate_rb/model_annotator/model_class_getter.rb

@@ -8,6 +8,13 @@ class << self
         # Check for namespaced models in subdirectories as well as models
         # in subdirectories without namespacing.
         def call(file, options)
+          use_zeitwerk = defined?(::Rails) && ::Rails.try(:autoloaders).try(:zeitwerk_enabled?)
+
+          if use_zeitwerk
+            klass = ZeitwerkClassGetter.call(file, options)
+            return klass if klass
+          end
+
           model_path = file.gsub(/\.rb$/, "")
           options[:model_dir].each { |dir| model_path = model_path.gsub(/^#{dir}/, "").gsub(/^\//, "") }
 

lib/annotate_rb/model_annotator/model_wrapper.rb

@@ -115,7 +115,16 @@ def retrieve_indexes_from_table
 
         # Try to search the table without prefix
         table_name_without_prefix = table_name.to_s.sub(@klass.table_name_prefix, "")
-        @klass.connection.indexes(table_name_without_prefix)
+        begin
+          @klass.connection.indexes(table_name_without_prefix)
+        rescue ActiveRecord::StatementInvalid => _e
+          # Mysql2 adapter behaves differently than Sqlite3 and Postgres adapter.
+          # If `table_name_without_prefix` does not exist, Mysql2 will raise,
+          # the other adapters will return an empty array.
+          #
+          # See: https://github.com/rails/rails/issues/51205
+          []
+        end
       end
 
       def with_comments?

lib/annotate_rb/model_annotator/zeitwerk_class_getter.rb

@@ -0,0 +1,113 @@
+# frozen_string_literal: true
+
+module AnnotateRb
+  module ModelAnnotator
+    class ZeitwerkClassGetter
+      class << self
+        def call(file, options)
+          new(file, options).call
+        end
+      end
+
+      def initialize(file, options)
+        @file = file
+        @options = options
+      end
+
+      # @return [Constant, nil] Attempts to return the model class constant (e.g. User) defined in the model file
+      #   can return `nil` if the file does not define the constant.
+      def call
+        return unless defined?(::Zeitwerk)
+
+        @absolute_file_path = File.expand_path(@file)
+        loader = ::Rails.autoloaders.main
+
+        if supports_cpath?
+          constant_using_cpath(loader)
+        else
+          constant(loader)
+        end
+      end
+
+      private
+
+      def constant(loader)
+        root_dirs = loader.dirs(namespaces: true) # or `root_dirs = loader.root_dirs` with zeitwerk < 2.6.1
+        expanded_file = @absolute_file_path
+
+        # root_dir: "/home/dummyapp/app/models"
+        root_dir, namespace = root_dirs.find do |dir, _namespace|
+          expanded_file.start_with?(dir)
+        end
+
+        # expanded_file: "/home/dummyapp/app/models/collapsed/example/test_model.rb"
+        # filepath_relative_to_root_dir: "/collapsed/example/test_model.rb"
+        _, filepath_relative_to_root_dir = expanded_file.split(root_dir)
+
+        # Remove leading / and the .rb extension
+        filepath_relative_to_root_dir = filepath_relative_to_root_dir[1..].sub(/\.rb$/, "")
+
+        # once we have the filepath_relative_to_root_dir, we need to see if it
+        # falls within one of our Zeitwerk "collapsed" paths.
+        if loader.collapse.any? { |path| path.include?(root_dir) && file.include?(path.split(root_dir)[1]) }
+          # if the file is within a collapsed path, we then need to, for each
+          # collapsed path, remove the root dir
+          collapsed = loader.collapse.map { |path| path.split(root_dir)[1].sub(/^\//, "") }.to_set
+
+          collapsed.each do |collapse|
+            # next, we split the collapsed directory, e.g. `domain_name/models`, by
+            # slash, and discard the domain_name
+            _, *collapsed_namespace = collapse.split("/")
+
+            # if there are any collapsed namespaces, e.g. `models`, we then remove
+            # that from `filepath_relative_to_root_dir`.
+            #
+            # This would result in:
+            #
+            # previous filepath_relative_to_root_dir: domain_name/models/model_name
+            # new filepath_relative_to_root_dir: domain_name/model_name
+            if collapsed_namespace.any?
+              filepath_relative_to_root_dir.sub!("/#{collapsed_namespace.last}", "")
+            end
+          end
+        end
+
+        camelize = loader.inflector.camelize(filepath_relative_to_root_dir, nil)
+        namespace.const_get(camelize)
+      rescue NameError => e
+        warn e
+        nil
+      end
+
+      def constant_using_cpath(loader)
+        begin
+          constant = loader.cpath_expected_at(@absolute_file_path)
+        rescue ::Zeitwerk::Error => e
+          # Raises when file does not exist
+          warn "Zeitwerk unable to find file #{@file}, error:\n#{e.message}"
+          return
+        end
+
+        begin
+          # This uses ActiveSupport::Inflector.constantize
+          klass = constant.constantize
+        rescue NameError => e
+          warn e
+          return
+        end
+
+        klass
+      end
+
+      def supports_cpath?
+        @supports_cpath ||=
+          begin
+            current_version = ::Gem::Version.new(::Zeitwerk::VERSION)
+            required_version = ::Gem::Version.new("2.6.9")
+
+            current_version >= required_version
+          end
+      end
+    end
+  end
+end

spec/dummyapp/app/models/collapsed/example/test_model.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+module Collapsed
+  class TestModel < ApplicationRecord
+    def self.table_name_prefix
+      "collapsed_"
+    end
+  end
+end

spec/dummyapp/config/initializers/test_zeitwork_collapsed.rb

@@ -0,0 +1,3 @@
+# frozen_string_literal: true
+
+Rails.autoloaders.main.collapse("#{Rails.root}/app/models/collapsed/example")

spec/dummyapp/db/migrate/20240210100506_create_collapsed_test_models.rb

@@ -0,0 +1,10 @@
+class CreateCollapsedTestModels < ActiveRecord::Migration[7.0]
+  def change
+    create_table :collapsed_test_models do |t|
+      t.string :name
+      t.boolean :collapsed
+
+      t.timestamps
+    end
+  end
+end

spec/integration/annotate_after_migration_spec.rb

@@ -22,7 +22,8 @@
 
     copy(File.join(migrations_template_dir, migration_file), "db/migrate")
 
-    _run_migrations_cmd = run_command_and_stop("bin/rails db:migrate VERSION=20231013230731", fail_on_error: true, exit_timeout: command_timeout_seconds)
+    # Apply this specific migration
+    _run_migrations_cmd = run_command_and_stop("bin/rails db:migrate:up VERSION=20231013230731", fail_on_error: true, exit_timeout: command_timeout_seconds)
     _run_annotations_cmd = run_command_and_stop("bundle exec annotaterb models", fail_on_error: true, exit_timeout: command_timeout_seconds)
 
     annotated_test_default = read_file(dummyapp_model("test_default.rb"))

spec/integration/annotate_collapsed_models_spec.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+require "integration_spec_helper"
+
+RSpec.describe "Annotate collapsed models", type: "aruba" do
+  let(:models_dir) { "app/models" }
+  let(:command_timeout_seconds) { 10 }
+
+  context "when annotating collapsed models" do
+    it "annotates them correctly" do
+      reset_database
+      run_migrations
+
+      expected_test_model = read_file(model_template("collapsed_test_model.rb"))
+
+      original_test_model = read_file(dummyapp_model("collapsed/example/test_model.rb"))
+
+      expect(expected_test_model).not_to eq(original_test_model)
+
+      _cmd = run_command_and_stop("bundle exec annotaterb models", fail_on_error: true, exit_timeout: command_timeout_seconds)
+
+      annotated_test_model = read_file(dummyapp_model("collapsed/example/test_model.rb"))
+
+      expect(last_command_started).to be_successfully_executed
+      expect(expected_test_model).to eq(annotated_test_model)
+    end
+  end
+end