Add polymorphic option to model generator

For instance,

    $ rails g model Product supplier:references{polymorphic}

generate model with `belongs_to :supplier, polymorphic: true` association and appropriate migration.

Also fix model_generator_test.rb#L196 and #L201

3 files changed, 41 insertions, 13 deletions

diff --git a/railties/lib/rails/generators/generated_attribute.rb b/railties/lib/rails/generators/generated_attribute.rb
index 25d0161e4c..d480fc12b5 100644
--- a/railties/lib/rails/generators/generated_attribute.rb
+++ b/railties/lib/rails/generators/generated_attribute.rb
@@ -35,7 +35,7 @@ module Rails
 
         private
 
-        # parse possible attribute options like :limit for string/text/binary/integer or :precision/:scale for decimals
+        # parse possible attribute options like :limit for string/text/binary/integer, :precision/:scale for decimals or :polymorphic for references/belongs_to
         # when declaring options curly brackets should be used
         def parse_type_and_options(type)
           case type
@@ -43,6 +43,8 @@ module Rails
             return $1, :limit => $2.to_i
           when /decimal\{(\d+)[,.-](\d+)\}/
             return :decimal, :precision => $1.to_i, :scale => $2.to_i
+          when /(references|belongs_to)\{polymorphic\}/
+            return $1, :polymorphic => true
           else
             return type, {}
           end
@@ -92,13 +94,21 @@ module Rails
       end
 
       def index_name
-        reference? ? "#{name}_id" : name
+        if reference?
+          polymorphic? ? %w(id type).map { |t| "#{name}_#{t}" } : "#{name}_id"
+        else
+          name
+        end
       end
 
       def reference?
         self.class.reference?(type)
       end
 
+      def polymorphic?
+        self.attr_options.has_key?(:polymorphic)
+      end
+
       def has_index?
         @has_index
       end
diff --git a/railties/test/generators/generated_attribute_test.rb b/railties/test/generators/generated_attribute_test.rb
index 6e3fc84781..6ab1cd58c7 100644
--- a/railties/test/generators/generated_attribute_test.rb
+++ b/railties/test/generators/generated_attribute_test.rb
@@ -102,22 +102,28 @@ class GeneratedAttributeTest < Rails::Generators::TestCase
 
   def test_reference_is_true
     %w(references belongs_to).each do |attribute_type|
-      assert_equal(
-        true,
-        create_generated_attribute(attribute_type).reference?
-      )
+      assert create_generated_attribute(attribute_type).reference?
     end
   end
 
   def test_reference_is_false
     %w(foo bar baz).each do |attribute_type|
-      assert_equal(
-        false,
-        create_generated_attribute(attribute_type).reference?
-      )
+      assert !create_generated_attribute(attribute_type).reference?
     end
   end
 
+  def test_polymorphic_reference_is_true
+    %w(references belongs_to).each do |attribute_type|
+      assert create_generated_attribute("#{attribute_type}{polymorphic}").polymorphic?
+    end
+  end
+  
+  def test_polymorphic_reference_is_false
+    %w(foo bar baz).each do |attribute_type|
+      assert !create_generated_attribute("#{attribute_type}{polymorphic}").polymorphic?
+    end
+  end
+  
   def test_blank_type_defaults_to_string_raises_exception
     assert_equal :string, create_generated_attribute(nil, 'title').type
     assert_equal :string, create_generated_attribute("", 'title').type
@@ -126,5 +132,6 @@ class GeneratedAttributeTest < Rails::Generators::TestCase
   def test_handles_index_names_for_references
     assert_equal "post", create_generated_attribute('string', 'post').index_name
     assert_equal "post_id", create_generated_attribute('references', 'post').index_name
+    assert_equal ["post_id", "post_type"], create_generated_attribute('references{polymorphic}', 'post').index_name
   end
 end
diff --git a/railties/test/generators/model_generator_test.rb b/railties/test/generators/model_generator_test.rb
index fd3b8c8a17..ec33bd7c6b 100644
--- a/railties/test/generators/model_generator_test.rb
+++ b/railties/test/generators/model_generator_test.rb
@@ -166,7 +166,7 @@ class ModelGeneratorTest < Rails::Generators::TestCase
   end
 
   def test_add_migration_with_attributes_index_declaration_and_attribute_options
-    run_generator ["product", "title:string{40}:index", "content:string{255}", "price:decimal{5,2}:index", "discount:decimal{5,2}:uniq"]
+    run_generator ["product", "title:string{40}:index", "content:string{255}", "price:decimal{5,2}:index", "discount:decimal{5,2}:uniq", "supplier:references{polymorphic}"]
 
     assert_migration "db/migrate/create_products.rb" do |content|
       assert_method :change, content do |up|
@@ -174,6 +174,7 @@ class ModelGeneratorTest < Rails::Generators::TestCase
         assert_match(/t.string :title, limit: 40/, up)
         assert_match(/t.string :content, limit: 255/, up)
         assert_match(/t.decimal :price, precision: 5, scale: 2/, up)
+        assert_match(/t.references :supplier, polymorphic: true/, up)
       end
       assert_match(/add_index :products, :title/, content)
       assert_match(/add_index :products, :price/, content)
@@ -193,15 +194,25 @@ class ModelGeneratorTest < Rails::Generators::TestCase
   end
 
   def test_model_with_references_attribute_generates_belongs_to_associations
-    run_generator ["product", "name:string", "supplier_id:references"]
+    run_generator ["product", "name:string", "supplier:references"]
     assert_file "app/models/product.rb", /belongs_to :supplier/
   end
 
   def test_model_with_belongs_to_attribute_generates_belongs_to_associations
-    run_generator ["product", "name:string", "supplier_id:belongs_to"]
+    run_generator ["product", "name:string", "supplier:belongs_to"]
     assert_file "app/models/product.rb", /belongs_to :supplier/
   end
 
+  def test_model_with_polymorphic_references_attribute_generates_belongs_to_associations
+    run_generator ["product", "name:string", "supplier:references{polymorphic}"]
+    assert_file "app/models/product.rb", /belongs_to :supplier, polymorphic: true/
+  end
+
+  def test_model_with_polymorphic_belongs_to_attribute_generates_belongs_to_associations
+    run_generator ["product", "name:string", "supplier:belongs_to{polymorphic}"]
+    assert_file "app/models/product.rb", /belongs_to :supplier, polymorphic: true/
+  end
+
   def test_migration_with_timestamps
     run_generator
     assert_migration "db/migrate/create_accounts.rb", /t.timestamps/