Optimize indentation for generator actions

4 files changed, 99 insertions, 19 deletions

diff --git a/railties/CHANGELOG.md b/railties/CHANGELOG.md
index f9d6258b22..408d0a07ff 100644
--- a/railties/CHANGELOG.md
+++ b/railties/CHANGELOG.md
@@ -1,3 +1,7 @@
+*   Optimize indentation for generator actions.
+
+    *Yoshiyuki Hirano*
+
 *   Add `--skip-action-cable` option to the plugin generator.
 
     *bogdanvlviv*
diff --git a/railties/lib/rails/generators.rb b/railties/lib/rails/generators.rb
index b9a0fe07c8..8915441d4c 100644
--- a/railties/lib/rails/generators.rb
+++ b/railties/lib/rails/generators.rb
@@ -10,6 +10,7 @@ require "active_support/core_ext/kernel/singleton_class"
 require "active_support/core_ext/array/extract_options"
 require "active_support/core_ext/hash/deep_merge"
 require "active_support/core_ext/module/attribute_accessors"
+require "active_support/core_ext/string/indent"
 require "active_support/core_ext/string/inflections"
 
 module Rails
diff --git a/railties/lib/rails/generators/actions.rb b/railties/lib/rails/generators/actions.rb
index 2792d7636f..9baf53c1d0 100644
--- a/railties/lib/rails/generators/actions.rb
+++ b/railties/lib/rails/generators/actions.rb
@@ -97,16 +97,16 @@ module Rails
       #     "config.action_controller.asset_host = 'localhost:3000'"
       #   end
       def environment(data = nil, options = {})
-        sentinel = /class [a-z_:]+ < Rails::Application/i
-        env_file_sentinel = /Rails\.application\.configure do/
-        data = yield if !data && block_given?
+        sentinel = "class Application < Rails::Application\n"
+        env_file_sentinel = "Rails.application.configure do\n"
+        data ||= yield if block_given?
 
         in_root do
           if options[:env].nil?
-            inject_into_file "config/application.rb", "\n    #{data}", after: sentinel, verbose: false
+            inject_into_file "config/application.rb", optimize_indentation(data, 4), after: sentinel, verbose: false
           else
             Array(options[:env]).each do |env|
-              inject_into_file "config/environments/#{env}.rb", "\n  #{data}", after: env_file_sentinel, verbose: false
+              inject_into_file "config/environments/#{env}.rb", optimize_indentation(data, 2), after: env_file_sentinel, verbose: false
             end
           end
         end
@@ -137,9 +137,10 @@ module Rails
       #   end
       #
       #   vendor("foreign.rb", "# Foreign code is fun")
-      def vendor(filename, data = nil, &block)
+      def vendor(filename, data = nil)
         log :vendor, filename
-        create_file("vendor/#{filename}", data, verbose: false, &block)
+        data ||= yield if block_given?
+        create_file("vendor/#{filename}", optimize_indentation(data), verbose: false)
       end
 
       # Create a new file in the lib/ directory. Code can be specified
@@ -150,9 +151,10 @@ module Rails
       #   end
       #
       #   lib("foreign.rb", "# Foreign code is fun")
-      def lib(filename, data = nil, &block)
+      def lib(filename, data = nil)
         log :lib, filename
-        create_file("lib/#{filename}", data, verbose: false, &block)
+        data ||= yield if block_given?
+        create_file("lib/#{filename}", optimize_indentation(data), verbose: false)
       end
 
       # Create a new +Rakefile+ with the provided code (either in a block or a string).
@@ -170,9 +172,10 @@ module Rails
       #   end
       #
       #   rakefile('seed.rake', 'puts "Planting seeds"')
-      def rakefile(filename, data = nil, &block)
+      def rakefile(filename, data = nil)
         log :rakefile, filename
-        create_file("lib/tasks/#{filename}", data, verbose: false, &block)
+        data ||= yield if block_given?
+        create_file("lib/tasks/#{filename}", optimize_indentation(data), verbose: false)
       end
 
       # Create a new initializer with the provided code (either in a block or a string).
@@ -188,9 +191,10 @@ module Rails
       #   end
       #
       #   initializer("api.rb", "API_KEY = '123456'")
-      def initializer(filename, data = nil, &block)
+      def initializer(filename, data = nil)
         log :initializer, filename
-        create_file("config/initializers/#{filename}", data, verbose: false, &block)
+        data ||= yield if block_given?
+        create_file("config/initializers/#{filename}", optimize_indentation(data), verbose: false)
       end
 
       # Generate something using a generator from Rails or a plugin.
@@ -304,6 +308,17 @@ module Rails
             "'#{value}'"
           end
         end
+
+        # Returns optimized string with indentation
+        def optimize_indentation(value, amount = 0) # :doc:
+          return "#{value}\n" unless value.is_a?(String)
+
+          if value.lines.size > 1
+            value.strip_heredoc.indent(amount)
+          else
+            "#{value.strip.indent(amount)}\n"
+          end
+        end
     end
   end
 end
diff --git a/railties/test/generators/actions_test.rb b/railties/test/generators/actions_test.rb
index 03b29be907..d65ced55b8 100644
--- a/railties/test/generators/actions_test.rb
+++ b/railties/test/generators/actions_test.rb
@@ -139,14 +139,14 @@ class ActionsTest < Rails::Generators::TestCase
     run_generator
     autoload_paths = 'config.autoload_paths += %w["#{Rails.root}/app/extras"]'
     action :environment, autoload_paths
-    assert_file "config/application.rb", /  class Application < Rails::Application\n    #{Regexp.escape(autoload_paths)}/
+    assert_file "config/application.rb", /  class Application < Rails::Application\n    #{Regexp.escape(autoload_paths)}\n/
   end
 
   def test_environment_should_include_data_in_environment_initializer_block_with_env_option
     run_generator
     autoload_paths = 'config.autoload_paths += %w["#{Rails.root}/app/extras"]'
     action :environment, autoload_paths, env: "development"
-    assert_file "config/environments/development.rb", /Rails\.application\.configure do\n  #{Regexp.escape(autoload_paths)}/
+    assert_file "config/environments/development.rb", /Rails\.application\.configure do\n  #{Regexp.escape(autoload_paths)}\n/
   end
 
   def test_environment_with_block_should_include_block_contents_in_environment_initializer_block
@@ -163,6 +163,26 @@ class ActionsTest < Rails::Generators::TestCase
     end
   end
 
+  def test_environment_with_block_should_include_block_contents_with_multiline_data_in_environment_initializer_block
+    run_generator
+    data = <<-RUBY
+      config.encoding = "utf-8"
+      config.time_zone = "UTC"
+    RUBY
+    action(:environment) { data }
+    assert_file "config/application.rb", /  class Application < Rails::Application\n#{Regexp.escape(data.strip_heredoc.indent(4))}/
+  end
+
+  def test_environment_should_include_block_contents_with_multiline_data_in_environment_initializer_block_with_env_option
+    run_generator
+    data = <<-RUBY
+      config.encoding = "utf-8"
+      config.time_zone = "UTC"
+    RUBY
+    action(:environment, nil, env: "development") { data }
+    assert_file "config/environments/development.rb", /Rails\.application\.configure do\n#{Regexp.escape(data.strip_heredoc.indent(2))}/
+  end
+
   def test_git_with_symbol_should_run_command_using_git_scm
     assert_called_with(generator, :run, ["git init"]) do
       action :git, :init
@@ -177,22 +197,62 @@ class ActionsTest < Rails::Generators::TestCase
 
   def test_vendor_should_write_data_to_file_in_vendor
     action :vendor, "vendor_file.rb", "# vendor data"
-    assert_file "vendor/vendor_file.rb", "# vendor data"
+    assert_file "vendor/vendor_file.rb", "# vendor data\n"
+  end
+
+  def test_vendor_should_write_data_to_file_with_block_in_vendor
+    code = <<-RUBY
+      puts "one"
+      puts "two"
+      puts "three"
+    RUBY
+    action(:vendor, "vendor_file.rb") { code }
+    assert_file "vendor/vendor_file.rb", code.strip_heredoc
   end
 
   def test_lib_should_write_data_to_file_in_lib
     action :lib, "my_library.rb", "class MyLibrary"
-    assert_file "lib/my_library.rb", "class MyLibrary"
+    assert_file "lib/my_library.rb", "class MyLibrary\n"
+  end
+
+  def test_lib_should_write_data_to_file_with_block_in_lib
+    code = <<-RUBY
+      class MyLib
+        MY_CONSTANT = 123
+      end
+    RUBY
+    action(:lib, "my_library.rb") { code }
+    assert_file "lib/my_library.rb", code.strip_heredoc
   end
 
   def test_rakefile_should_write_date_to_file_in_lib_tasks
     action :rakefile, "myapp.rake", "task run: [:environment]"
-    assert_file "lib/tasks/myapp.rake", "task run: [:environment]"
+    assert_file "lib/tasks/myapp.rake", "task run: [:environment]\n"
+  end
+
+  def test_rakefile_should_write_date_to_file_with_block_in_lib_tasks
+    code = <<-RUBY
+      task rock: :environment do
+        puts "Rockin'"
+      end
+    RUBY
+    action(:rakefile, "myapp.rake") { code }
+    assert_file "lib/tasks/myapp.rake", code.strip_heredoc
   end
 
   def test_initializer_should_write_date_to_file_in_config_initializers
     action :initializer, "constants.rb", "MY_CONSTANT = 42"
-    assert_file "config/initializers/constants.rb", "MY_CONSTANT = 42"
+    assert_file "config/initializers/constants.rb", "MY_CONSTANT = 42\n"
+  end
+
+  def test_initializer_should_write_date_to_file_with_block_in_config_initializers
+    code = <<-RUBY
+      MyLib.configure do |config|
+        config.value = 123
+      end
+    RUBY
+    action(:initializer, "constants.rb") { code }
+    assert_file "config/initializers/constants.rb", code.strip_heredoc
   end
 
   def test_generate_should_run_script_generate_with_argument_and_options