Select from all possible load paths

lib/active_record/annotate.rb

@@ -9,67 +9,74 @@ module Annotate
     class << self
       def annotate
         processed_models = []
-        
+
         models.each do |table_name, file_paths_and_classes|
           annotation = Dumper.dump(table_name)
-          
+
           file_paths_and_classes.each do |path, klass|
             file = File.new(path)
             file.annotate_with(annotation.dup, configurator)
-            
+
             if file.changed?
               file.write
               processed_models << "#{klass} (#{file.relative_path})"
             end
           end
         end
-        
+
         unless processed_models.empty?
           puts 'Annotated models:'
           processed_models.each do |model|
             puts "  * #{model}"
           end
         end
       end
-      
+
       def models
-        files_mask = models_dir.join('**', '*.rb')
-
+        files_mask = models_dirs.map do |path|
+          ::File.join(path, '**', '*.rb')
+        end
+
         hash_with_arrays = Hash.new do |hash, key|
           hash[key] = []
         end
-        
+
         Dir.glob(files_mask).each_with_object(hash_with_arrays) do |path, models|
           short_path = short_path_for(path)
           next if short_path.starts_with?('concerns') # skip any app/models/concerns files
-          
+
           klass = class_name_for(short_path)
           next unless klass < ActiveRecord::Base # collect only AR::Base descendants
           next if klass.respond_to?(:abstract_class?) && klass.abstract_class?
-          
+
           models[klass.table_name] << [path, klass]
         end
       end
-      
+
       # .../app/models/car/hatchback.rb -> car/hatchback
       def short_path_for(full_path)
+        models_dir = models_dirs.find { |path| full_path.include?(path) }
+
         full_path.sub(models_dir.to_s + '/', '').sub(/\.rb$/, '')
       end
-      
+
       # car/hatchback -> Car::Hatchback
       def class_name_for(short_path)
         short_path.camelize.constantize
       end
-      
+
       def configure(&block)
         configurator.tap(&block)
       end
-
-    private
-      def models_dir
-        Rails.root.join('app/models')
+
+      private
+
+      def models_dirs
+        Rails.application.paths.all_paths.flat_map(&:to_a).select do |path|
+          path.end_with?('models')
+        end
       end
-      
+
       def configurator
         @configurator ||= Configurator.new
       end

spec/active_record/annotate_spec.rb

@@ -1,19 +1,19 @@
 require 'spec_helper'
 
 describe ActiveRecord::Annotate do
-  describe ".short_path_for" do
+  describe '.short_path_for' do
     before(:each) do
-      allow(subject).to receive(:models_dir).and_return('dir')
+      allow(subject).to receive(:models_dirs).and_return(['dir'])
     end
-    
-    it "removes the root/app/models prefix and .rb suffix" do
+
+    it 'removes the root/app/models prefix and .rb suffix' do
       short_path = subject.short_path_for('dir/namespace/model_name.rb')
       expect(short_path).to eq('namespace/model_name')
     end
   end
-  
-  describe ".class_name_for" do
-    it "finds the class by the short path" do
+
+  describe '.class_name_for' do
+    it 'finds the class by the short path' do
       class_name = subject.class_name_for('active_record/annotate/file')
       expect(class_name).to eq(ActiveRecord::Annotate::File)
     end