Refactor application of type decorators into an object

1 files changed, 24 insertions, 8 deletions

diff --git a/activerecord/lib/active_record/attribute_decorators.rb b/activerecord/lib/active_record/attribute_decorators.rb
index 596161f81d..928d8cacae 100644
--- a/activerecord/lib/active_record/attribute_decorators.rb
+++ b/activerecord/lib/active_record/attribute_decorators.rb
@@ -4,7 +4,7 @@ module ActiveRecord
 
     included do
       class_attribute :attribute_type_decorations, instance_accessor: false # :internal:
-      self.attribute_type_decorations = Hash.new({})
+      self.attribute_type_decorations = TypeDecorator.new
     end
 
     module ClassMethods
@@ -13,22 +13,38 @@ module ActiveRecord
         column_name = column_name.to_s
 
         # Create new hashes so we don't modify parent classes
-        decorations_for_column = attribute_type_decorations[column_name]
-        new_decorations = decorations_for_column.merge(decorator_name.to_s => block)
-        self.attribute_type_decorations = attribute_type_decorations.merge(column_name => new_decorations)
+        self.attribute_type_decorations = attribute_type_decorations.merge(column_name, decorator_name, block)
       end
 
       private
 
       def add_user_provided_columns(*)
         super.map do |column|
-          decorations = attribute_type_decorations[column.name].values
-          decorated_type = decorations.inject(column.cast_type) do |type, block|
-            block.call(type)
-          end
+          decorated_type = attribute_type_decorations.apply(column.name, column.cast_type)
           column.with_type(decorated_type)
         end
       end
     end
+
+    class TypeDecorator
+      delegate :clear, to: :@decorations
+
+      def initialize(decorations = Hash.new({}))
+        @decorations = decorations
+      end
+
+      def merge(attribute_name, decorator_name, block)
+        decorations_for_attribute = @decorations[attribute_name]
+        new_decorations = decorations_for_attribute.merge(decorator_name.to_s => block)
+        TypeDecorator.new(@decorations.merge(attribute_name => new_decorations))
+      end
+
+      def apply(attribute_name, type)
+        decorations = @decorations[attribute_name].values
+        decorations.inject(type) do |new_type, block|
+          block.call(new_type)
+        end
+      end
+    end
   end
 end