Move ActiveStorage reflection logic entirely into ActiveStorage

2 files changed, 72 insertions, 0 deletions

diff --git a/activestorage/lib/active_storage/engine.rb b/activestorage/lib/active_storage/engine.rb
index 99588cdd4b..519b9ae283 100644
--- a/activestorage/lib/active_storage/engine.rb
+++ b/activestorage/lib/active_storage/engine.rb
@@ -10,6 +10,8 @@ require "active_storage/previewer/video_previewer"
 require "active_storage/analyzer/image_analyzer"
 require "active_storage/analyzer/video_analyzer"
 
+require "active_storage/reflection"
+
 module ActiveStorage
   class Engine < Rails::Engine # :nodoc:
     isolate_namespace ActiveStorage
@@ -95,5 +97,12 @@ module ActiveStorage
         end
       end
     end
+
+    initializer "active_storage.reflection" do
+      ActiveSupport.on_load(:active_record) do
+        include Reflection::ActiveRecordExtensions
+        ActiveRecord::Reflection.singleton_class.prepend(Reflection::ReflectionExtension)
+      end
+    end
   end
 end
diff --git a/activestorage/lib/active_storage/reflection.rb b/activestorage/lib/active_storage/reflection.rb
new file mode 100644
index 0000000000..9074b20126
--- /dev/null
+++ b/activestorage/lib/active_storage/reflection.rb
@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+module ActiveStorage
+  module Reflection
+    # Holds all the metadata about a has_one_attached attachment as it was
+    # specified in the Active Record class.
+    class HasOneAttachedReflection < ActiveRecord::Reflection::MacroReflection #:nodoc:
+      def macro
+        :has_one_attached
+      end
+    end
+
+    # Holds all the metadata about a has_many_attached attachment as it was
+    # specified in the Active Record class.
+    class HasManyAttachedReflection < ActiveRecord::Reflection::MacroReflection #:nodoc:
+      def macro
+        :has_many_attached
+      end
+    end
+
+    module ReflectionExtension
+      def reflection_class_for(macro)
+        case macro
+        when :has_one_attached
+          HasOneAttachedReflection
+        when :has_many_attached
+          HasManyAttachedReflection
+        else
+          super
+        end
+      end
+
+      def add_attachment_reflection(ar, name, reflection)
+        ar.attachment_reflections.merge!(name.to_s => reflection)
+      end
+    end
+
+    module ActiveRecordExtensions
+      extend ActiveSupport::Concern
+
+      included do
+        class_attribute :attachment_reflections, instance_writer: false, default: {}
+      end
+
+      module ClassMethods
+        # Returns an array of reflection objects for all the attachments in the
+        # class.
+        def reflect_on_all_attachments
+          attachment_reflections.values
+        end
+
+        # Returns the reflection object for the named +attachment+.
+        #
+        #    User.reflect_on_attachment(:avatar)
+        #    # => the avatar reflection
+        #
+        def reflect_on_attachment(attachment)
+          attachment_reflections[attachment.to_s]
+        end
+      end
+    end
+  end
+end