Merge pull request rails#50796 from Shopify/implement-#50770

Add customized prompt for Rails console

Author: byroot

railties/CHANGELOG.md

@@ -1,3 +1,18 @@
+*   Rails console now indicates the current Rails environment:
+
+    ```txt
+    dev:001> # for RAILS_ENV=development
+    test:001> # for RAILS_ENV=test
+    prod:001> # for RAILS_ENV=production
+    my_env:001> # for RAILS_ENV=my_env
+    ```
+
+    The environment name will also be colorized when the environment is
+    `development` (green), `test` (green), or `production` (red), if your
+    terminal supports it.
+
+    *Stan Lo*
+
 *   `bin/rails` now prints its help message when given an unrecognized bare
     option.
 

railties/lib/rails/commands/console/console_command.rb

@@ -12,6 +12,53 @@ def filter_backtrace(bt)
       end
     end
 
+    class IRBConsole
+      def initialize
+        require "irb"
+        require "irb/completion"
+        IRB::WorkSpace.prepend(BacktraceCleaner)
+
+        if !Rails.env.local?
+          # Use env var here so users can override them with env var too
+          ENV["IRB_USE_AUTOCOMPLETE"] ||= "false"
+        end
+      end
+
+      def name
+        "IRB"
+      end
+
+      def start
+        IRB.setup(nil)
+
+        env = colorized_env
+
+        IRB.conf[:PROMPT][:RAILS_PROMPT] = {
+          PROMPT_I: "#{env}:%03n> ",
+          PROMPT_S: "#{env}:%03n%l ",
+          PROMPT_C: "#{env}:%03n* ",
+          RETURN: "=> %s\n"
+        }
+
+        # Respect user's choice of prompt mode.
+        IRB.conf[:PROMPT_MODE] = :RAILS_PROMPT if IRB.conf[:PROMPT_MODE] == :DEFAULT
+        IRB::Irb.new.run(IRB.conf)
+      end
+
+      def colorized_env
+        case Rails.env
+        when "development"
+          IRB::Color.colorize("dev", [:GREEN])
+        when "test"
+          IRB::Color.colorize("test", [:GREEN])
+        when "production"
+          IRB::Color.colorize("prod", [:RED])
+        else
+          Rails.env
+        end
+      end
+    end
+
     def self.start(*args)
       new(*args).start
     end
@@ -31,18 +78,7 @@ def initialize(app, options = {})
 
       app.load_console
 
-      @console = app.config.console || begin
-        require "irb"
-        require "irb/completion"
-
-        IRB::WorkSpace.prepend(BacktraceCleaner)
-
-        if !Rails.env.local?
-          ENV["IRB_USE_AUTOCOMPLETE"] ||= "false"
-        end
-
-        IRB
-      end
+      @console = app.config.console || IRBConsole.new
     end
 
     def sandbox?

railties/test/application/console_test.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require "tempfile"
+
 require "isolation/abstract_unit"
 require "console_helpers"
 
@@ -123,9 +125,9 @@ def write_prompt(command, expected_output = nil)
     assert_output "> ", @primary
   end
 
-  def spawn_console(options, wait_for_prompt: true)
+  def spawn_console(options, wait_for_prompt: true, env: {})
     pid = Process.spawn(
-      { "TERM" => "dumb" },
+      { "TERM" => "dumb" }.merge(env),
       "#{app_path}/bin/rails console #{options}",
       in: @replica, out: @replica, err: @replica
     )
@@ -209,4 +211,40 @@ def test_environment_option_and_irb_option
     write_prompt "puts Rails.env", "puts Rails.env\r\ntest"
     @primary.puts "quit"
   end
+
+  def test_production_console_prompt
+    options = "-e production -- --nocolorize"
+    spawn_console(options)
+
+    write_prompt "123", "prod:001> 123"
+  end
+
+  def test_development_console_prompt
+    options = "-e development -- --nocolorize"
+    spawn_console(options)
+
+    write_prompt "123", "dev:001> 123"
+  end
+
+  def test_test_console_prompt
+    options = "-e test -- --nocolorize"
+    spawn_console(options)
+
+    write_prompt "123", "test:001> 123"
+  end
+
+  def test_console_respects_user_defined_prompt_mode
+    irbrc = Tempfile.new("irbrc")
+    irbrc.write <<-RUBY
+      IRB.conf[:PROMPT_MODE] = :SIMPLE
+    RUBY
+    irbrc.close
+
+    options = "-e test -- --nocolorize"
+    spawn_console(options, env: { "IRBRC" => irbrc.path })
+
+    write_prompt "123", ">> 123"
+  ensure
+    File.unlink(irbrc)
+  end
 end

railties/test/commands/console_test.rb

@@ -99,6 +99,25 @@ def test_console_doesnt_disable_IRB_auto_completion_in_local
     ENV["IRB_USE_AUTOCOMPLETE"] = original_use_autocomplete
   end
 
+  def test_prompt_env_colorization
+    irb_console = Rails::Console::IRBConsole.new
+    red = "\e[31m"
+    green = "\e[32m"
+    clear = "\e[0m"
+
+    Rails.env = "development"
+    assert_equal("#{green}dev#{clear}", irb_console.colorized_env)
+
+    Rails.env = "test"
+    assert_equal("#{green}test#{clear}", irb_console.colorized_env)
+
+    Rails.env = "production"
+    assert_equal("#{red}prod#{clear}", irb_console.colorized_env)
+
+    Rails.env = "custom_env"
+    assert_equal("custom_env", irb_console.colorized_env)
+  end
+
   def test_default_environment_with_no_rails_env
     with_rails_env nil do
       start