Added STI support to init and building associations

Allows you to do BaseClass.new(:type => "SubClass") as well as
parent.children.build(:type => "SubClass") or parent.build_child
to initialize an STI subclass. Ensures that the class name is a
valid class and that it is in the ancestors of the super class
that the association is expecting.

1 files changed, 26 insertions, 0 deletions

diff --git a/activerecord/lib/active_record/inheritance.rb b/activerecord/lib/active_record/inheritance.rb
index a448fa1f5c..850911ebe7 100644
--- a/activerecord/lib/active_record/inheritance.rb
+++ b/activerecord/lib/active_record/inheritance.rb
@@ -9,6 +9,19 @@ module ActiveRecord
     end
 
     module ClassMethods
+      # Determines if one of the attributes passed in is the inheritance column,
+      # and if the inheritance column is attr accessible, it initializes an
+      # instance of the given subclass instead of the base class
+      def new(*args, &block)
+        if (attrs = args.first).is_a?(Hash)
+          if subclass = subclass_from_attrs(attrs)
+            return subclass.new(*args, &block)
+          end
+        end
+        # Delegate to the original .new
+        super
+      end
+
       # True if this isn't a concrete subclass needing a STI type condition.
       def descends_from_active_record?
         if self == Base
@@ -145,6 +158,19 @@ module ActiveRecord
 
         sti_column.in(sti_names)
       end
+
+      # Detect the subclass from the inheritance column of attrs. If the inheritance column value
+      # is not self or a valid subclass, raises ActiveRecord::SubclassNotFound
+      # If this is a StrongParameters hash, and access to inheritance_column is not permitted,
+      # this will ignore the inheritance column and return nil
+      def subclass_from_attrs(attrs)
+        subclass_name = attrs.with_indifferent_access[inheritance_column]
+        return nil if subclass_name.blank? || subclass_name == self.name
+        unless subclass = subclasses.detect { |sub| sub.name == subclass_name }
+          raise ActiveRecord::SubclassNotFound.new("Invalid single-table inheritance type: #{subclass_name} is not a subclass of #{name}")
+        end
+        subclass
+      end
     end
 
     private