Promote time zone aware attributes to a first class type decorator

This refactoring revealed the need for another form of decoration, which
takes a proc to select which it applies to (There's a *lot* of cases
where this form can be used). To avoid duplication, we can re-implement
the old decoration in terms of the proc-based decoration.

The reason we're `instance_exec`ing the matcher is for cases such as
time zone aware attributes, where a decorator is defined in a parent
class, and a method called in the matcher is overridden by a child
class. The matcher will close over the parent, and evaluate in its
context, which is not the behavior we want.

6 files changed, 57 insertions, 32 deletions

diff --git a/activerecord/lib/active_record/attribute_decorators.rb b/activerecord/lib/active_record/attribute_decorators.rb
index 928d8cacae..92627f8d5d 100644
--- a/activerecord/lib/active_record/attribute_decorators.rb
+++ b/activerecord/lib/active_record/attribute_decorators.rb
@@ -9,18 +9,24 @@ module ActiveRecord
 
     module ClassMethods
       def decorate_attribute_type(column_name, decorator_name, &block)
+        matcher = ->(name, _) { name == column_name.to_s }
+        key = "_#{column_name}_#{decorator_name}"
+        decorate_matching_attribute_types(matcher, key, &block)
+      end
+
+      def decorate_matching_attribute_types(matcher, decorator_name, &block)
         clear_caches_calculated_from_columns
-        column_name = column_name.to_s
+        decorator_name = decorator_name.to_s
 
         # Create new hashes so we don't modify parent classes
-        self.attribute_type_decorations = attribute_type_decorations.merge(column_name, decorator_name, block)
+        self.attribute_type_decorations = attribute_type_decorations.merge(decorator_name => [matcher, block])
       end
 
       private
 
       def add_user_provided_columns(*)
         super.map do |column|
-          decorated_type = attribute_type_decorations.apply(column.name, column.cast_type)
+          decorated_type = attribute_type_decorations.apply(self, column.name, column.cast_type)
           column.with_type(decorated_type)
         end
       end
@@ -29,22 +35,32 @@ module ActiveRecord
     class TypeDecorator
       delegate :clear, to: :@decorations
 
-      def initialize(decorations = Hash.new({}))
+      def initialize(decorations = {})
         @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))
+      def merge(*args)
+        TypeDecorator.new(@decorations.merge(*args))
       end
 
-      def apply(attribute_name, type)
-        decorations = @decorations[attribute_name].values
+      def apply(context, name, type)
+        decorations = decorators_for(context, name, type)
         decorations.inject(type) do |new_type, block|
           block.call(new_type)
         end
       end
+
+      private
+
+      def decorators_for(context, name, type)
+        matching(context, name, type).map(&:last)
+      end
+
+      def matching(context, name, type)
+        @decorations.values.select do |(matcher, _)|
+          context.instance_exec(name, type, &matcher)
+        end
+      end
     end
   end
 end
diff --git a/activerecord/lib/active_record/attribute_methods/time_zone_conversion.rb b/activerecord/lib/active_record/attribute_methods/time_zone_conversion.rb
index abad949ef4..188d5d2aab 100644
--- a/activerecord/lib/active_record/attribute_methods/time_zone_conversion.rb
+++ b/activerecord/lib/active_record/attribute_methods/time_zone_conversion.rb
@@ -1,7 +1,7 @@
 module ActiveRecord
   module AttributeMethods
     module TimeZoneConversion
-      class Type < SimpleDelegator # :nodoc:
+      class TimeZoneConverter < SimpleDelegator # :nodoc:
         def type_cast_from_database(value)
           convert_time_to_time_zone(super)
         end
@@ -33,6 +33,11 @@ module ActiveRecord
 
         class_attribute :skip_time_zone_conversion_for_attributes, instance_writer: false
         self.skip_time_zone_conversion_for_attributes = []
+
+        matcher = ->(name, type) { create_time_zone_conversion_attribute?(name, type) }
+        decorate_matching_attribute_types(matcher, :_time_zone_conversion) do |type|
+          TimeZoneConverter.new(type)
+        end
       end
 
       module ClassMethods
diff --git a/activerecord/lib/active_record/base.rb b/activerecord/lib/active_record/base.rb
index e4d0abb8ef..0b788ea1f9 100644
--- a/activerecord/lib/active_record/base.rb
+++ b/activerecord/lib/active_record/base.rb
@@ -310,6 +310,8 @@ module ActiveRecord #:nodoc:
     include CounterCache
     include Locking::Optimistic
     include Locking::Pessimistic
+    include Attributes
+    include AttributeDecorators
     include AttributeMethods
     include Callbacks
     include Timestamp
@@ -323,8 +325,6 @@ module ActiveRecord #:nodoc:
     include Reflection
     include Serialization
     include Store
-    include Attributes
-    include AttributeDecorators
   end
 
   ActiveSupport.run_load_hooks(:active_record, Base)
diff --git a/activerecord/lib/active_record/model_schema.rb b/activerecord/lib/active_record/model_schema.rb
index 9e1afd32e6..b9558b0752 100644
--- a/activerecord/lib/active_record/model_schema.rb
+++ b/activerecord/lib/active_record/model_schema.rb
@@ -220,27 +220,13 @@ module ActiveRecord
       end
 
       def column_types # :nodoc:
-        @column_types ||= decorate_types(build_types_hash)
+        @column_types ||= Hash[columns.map { |column| [column.name, column.cast_type] }]
       end
 
       def type_for_attribute(attr_name) # :nodoc:
         column_types.fetch(attr_name) { Type::Value.new }
       end
 
-      def decorate_types(types) # :nodoc:
-        return if types.empty?
-
-        @time_zone_column_names ||= self.columns_hash.find_all do |name, col|
-          create_time_zone_conversion_attribute?(name, col)
-        end.map!(&:first)
-
-        @time_zone_column_names.each do |name|
-          types[name] = AttributeMethods::TimeZoneConversion::Type.new(types[name])
-        end
-
-        types
-      end
-
       # Returns a hash where the keys are column names and the values are
       # default values when instantiating the AR object for this table.
       def column_defaults
@@ -335,10 +321,6 @@ module ActiveRecord
           base.table_name
         end
       end
-
-      def build_types_hash
-        Hash[columns.map { |column| [column.name, column.cast_type] }]
-      end
     end
   end
 end
diff --git a/activerecord/test/cases/attribute_decorators_test.rb b/activerecord/test/cases/attribute_decorators_test.rb
index 35393753a2..bc3e9a8cf5 100644
--- a/activerecord/test/cases/attribute_decorators_test.rb
+++ b/activerecord/test/cases/attribute_decorators_test.rb
@@ -110,5 +110,26 @@ module ActiveRecord
 
       assert_equal 'whatever decorated!', column.default
     end
+
+    class Multiplier < SimpleDelegator
+      def type_cast_from_user(value)
+        return if value.nil?
+        value * 2
+      end
+      alias type_cast_from_database type_cast_from_user
+    end
+
+    test "decorating with a proc" do
+      Model.attribute :an_int, Type::Integer.new
+      type_is_integer = proc { |_, type| type.type == :integer }
+      Model.decorate_matching_attribute_types type_is_integer, :multiplier do |type|
+        Multiplier.new(type)
+      end
+
+      model = Model.new(a_string: 'whatever', an_int: 1)
+
+      assert_equal 'whatever', model.a_string
+      assert_equal 2, model.an_int
+    end
   end
 end
diff --git a/activerecord/test/cases/attribute_methods/read_test.rb b/activerecord/test/cases/attribute_methods/read_test.rb
index c0659fddef..4741ee8799 100644
--- a/activerecord/test/cases/attribute_methods/read_test.rb
+++ b/activerecord/test/cases/attribute_methods/read_test.rb
@@ -12,6 +12,7 @@ module ActiveRecord
         @klass = Class.new do
           def self.superclass; Base; end
           def self.base_class; self; end
+          def self.decorate_matching_attribute_types(*); end
 
           include ActiveRecord::AttributeMethods