Adding option to skip loading models from subdirectories

lib/annotate/annotate_models.rb

@@ -67,6 +67,19 @@ def root_dir
 
     attr_writer :root_dir
 
+    def skip_subdirectory_model_load
+      # This option is set in options[:skip_subdirectory_model_load]
+      # and stops the get_loaded_model method from loading a model from a subdir
+
+      if @skip_subdirectory_model_load.blank?
+        false
+      else
+        @skip_subdirectory_model_load
+      end
+    end
+
+    attr_writer :skip_subdirectory_model_load
+
     def get_patterns(options, pattern_types = [])
       current_patterns = []
       root_dir.each do |root_directory|
@@ -586,8 +599,10 @@ def get_model_class(file)
 
     # Retrieve loaded model class
     def get_loaded_model(model_path, file)
-      loaded_model_class = get_loaded_model_by_path(model_path)
-      return loaded_model_class if loaded_model_class
+      unless skip_subdirectory_model_load
+        loaded_model_class = get_loaded_model_by_path(model_path)
+        return loaded_model_class if loaded_model_class
+      end
 
       # We cannot get loaded model when `model_path` is loaded by Rails
       # auto_load/eager_load paths. Try all possible model paths one by one.
@@ -616,6 +631,7 @@ def get_loaded_model_by_path(model_path)
     def parse_options(options = {})
       self.model_dir = split_model_dir(options[:model_dir]) if options[:model_dir]
       self.root_dir = options[:root_dir] if options[:root_dir]
+      self.skip_subdirectory_model_load = options[:skip_subdirectory_model_load].present?
     end
 
     def split_model_dir(option_value)

spec/lib/annotate/annotate_models_spec.rb

@@ -147,7 +147,8 @@ def mock_column(name, type, options = {})
     let(:options) do
       {
         root_dir: '/root',
-        model_dir: 'app/models,app/one,  app/two   ,,app/three'
+        model_dir: 'app/models,app/one,  app/two   ,,app/three',
+        skip_subdirectory_model_load: false
       }
     end
 
@@ -174,6 +175,40 @@ def mock_column(name, type, options = {})
         is_expected.to eq(['app/models', 'app/one', 'app/two', 'app/three'])
       end
     end
+
+    describe '@skip_subdirectory_model_load' do
+      subject do
+        AnnotateModels.instance_variable_get(:@skip_subdirectory_model_load)
+      end
+
+      context 'option is set to true' do
+        let(:options) do
+          {
+            root_dir: '/root',
+            model_dir: 'app/models,app/one,  app/two   ,,app/three',
+            skip_subdirectory_model_load: true
+          }
+        end
+
+        it 'sets skip_subdirectory_model_load to true' do
+          is_expected.to eq(true)
+        end
+      end
+
+      context 'option is set to false' do
+        let(:options) do
+          {
+            root_dir: '/root',
+            model_dir: 'app/models,app/one,  app/two   ,,app/three',
+            skip_subdirectory_model_load: false
+          }
+        end
+
+        it 'sets skip_subdirectory_model_load to false' do
+          is_expected.to eq(false)
+        end
+      end
+    end
   end
 
   describe '.get_schema_info' do
@@ -2087,6 +2122,21 @@ class Bar::Foo
           expect(klass.name).to eq('Foo')
           expect(klass_2.name).to eq('Bar::Foo')
         end
+
+        it 'attempts to load the model path without expanding if skip_subdirectory_model_load is false' do
+          allow(AnnotateModels).to receive(:skip_subdirectory_model_load).and_return(false)
+          full_path = File.join(AnnotateModels.model_dir[0], filename_2)
+          expect(File).to_not receive(:expand_path).with(full_path)
+          AnnotateModels.get_model_class(full_path)
+        end
+
+        it 'does not attempt to load the model path without expanding if skip_subdirectory_model_load is true' do
+          $LOAD_PATH.unshift(AnnotateModels.model_dir[0])
+          allow(AnnotateModels).to receive(:skip_subdirectory_model_load).and_return(true)
+          full_path = File.join(AnnotateModels.model_dir[0], filename_2)
+          expect(File).to receive(:expand_path).with(full_path).and_call_original
+          AnnotateModels.get_model_class(full_path)
+        end
       end
 
       context 'one of the classes is nested in another class' do