From a13836e70d9cc2eb569911030cbd735d68b4042c Mon Sep 17 00:00:00 2001
From: Takashi Kokubun <takashikkbn@gmail.com>
Date: Thu, 3 Nov 2022 15:09:51 -0700
Subject: [PATCH] [ruby/irb] Allow non-identifier aliases like Pry's @ and $
 (https://github.com/ruby/irb/pull/426)

* Allow non-identifier aliases

* Move the configuration to IRB.conf

* Avoid abusing method lookup for symbol aliases

* Add more alias tests

* A small optimization

* Assume non-nil Context

* Load IRB.conf earlier

https://github.com/ruby/irb/commit/e23db5132e
---
 lib/irb.rb           |  3 +++
 lib/irb/context.rb   | 18 ++++++++++++++
 lib/irb/init.rb      |  2 ++
 lib/irb/ruby-lex.rb  |  6 +++++
 test/irb/test_cmd.rb | 58 ++++++++++++++++++++++++++++++++++++++++++++
 5 files changed, 87 insertions(+)

diff --git a/lib/irb.rb b/lib/irb.rb
index 749f3ee167..57ec9ebaeb 100644
--- a/lib/irb.rb
+++ b/lib/irb.rb
@@ -426,6 +426,9 @@ module IRB
     def initialize(workspace = nil, input_method = nil)
       @context = Context.new(self, workspace, input_method)
       @context.main.extend ExtendCommandBundle
+      @context.command_aliases.each do |alias_name, cmd_name|
+        @context.main.install_alias_method(alias_name, cmd_name)
+      end
       @signal_status = :IN_IRB
       @scanner = RubyLex.new
     end
diff --git a/lib/irb/context.rb b/lib/irb/context.rb
index d238da9350..d1ae2cb605 100644
--- a/lib/irb/context.rb
+++ b/lib/irb/context.rb
@@ -149,6 +149,8 @@ module IRB
       if @newline_before_multiline_output.nil?
         @newline_before_multiline_output = true
       end
+
+      @command_aliases = IRB.conf[:COMMAND_ALIASES]
     end
 
     # The top-level workspace, see WorkSpace#main
@@ -326,6 +328,9 @@ module IRB
     # See IRB@Command+line+options for more command line options.
     attr_accessor :back_trace_limit
 
+    # User-defined IRB command aliases
+    attr_accessor :command_aliases
+
     # Alias for #use_multiline
     alias use_multiline? use_multiline
     # Alias for #use_singleline
@@ -477,6 +482,13 @@ module IRB
         line = "begin ::Kernel.raise _; rescue _.class\n#{line}\n""end"
         @workspace.local_variable_set(:_, exception)
       end
+
+      # Transform a non-identifier alias (ex: @, $)
+      command = line.split(/\s/, 2).first
+      if original = symbol_alias(command)
+        line = line.gsub(/\A#{Regexp.escape(command)}/, original.to_s)
+      end
+
       set_last_value(@workspace.evaluate(self, line, irb_path, line_no))
     end
 
@@ -522,5 +534,11 @@ module IRB
     def local_variables # :nodoc:
       workspace.binding.local_variables
     end
+
+    # Return a command name if it's aliased from the argument and it's not an identifier.
+    def symbol_alias(command)
+      return nil if command.match?(/\A\w+\z/)
+      command_aliases[command.to_sym]
+    end
   end
 end
diff --git a/lib/irb/init.rb b/lib/irb/init.rb
index 5409528fae..09099f88b7 100644
--- a/lib/irb/init.rb
+++ b/lib/irb/init.rb
@@ -158,6 +158,8 @@ module IRB # :nodoc:
     @CONF[:LC_MESSAGES] = Locale.new
 
     @CONF[:AT_EXIT] = []
+
+    @CONF[:COMMAND_ALIASES] = {}
   end
 
   def IRB.set_measure_callback(type = nil, arg = nil, &block)
diff --git a/lib/irb/ruby-lex.rb b/lib/irb/ruby-lex.rb
index 544392228e..28029bbf4c 100644
--- a/lib/irb/ruby-lex.rb
+++ b/lib/irb/ruby-lex.rb
@@ -65,6 +65,12 @@ class RubyLex
             false
           end
         else
+          # Accept any single-line input starting with a non-identifier alias (ex: @, $)
+          command = code.split(/\s/, 2).first
+          if context.symbol_alias(command)
+            next true
+          end
+
           code.gsub!(/\s*\z/, '').concat("\n")
           ltype, indent, continue, code_block_open = check_state(code, context: context)
           if ltype or indent > 0 or continue or code_block_open
diff --git a/test/irb/test_cmd.rb b/test/irb/test_cmd.rb
index 060f70c9cc..034d825bbc 100644
--- a/test/irb/test_cmd.rb
+++ b/test/irb/test_cmd.rb
@@ -570,6 +570,24 @@ module TestIRB
       assert_match(%r[/irb\.rb], out)
     end
 
+    def test_show_source_alias
+      input = TestInputMethod.new([
+        "$ 'IRB.conf'\n",
+      ])
+      IRB.init_config(nil)
+      IRB.conf[:COMMAND_ALIASES] = { :'$' => :show_source }
+      workspace = IRB::WorkSpace.new(Object.new)
+      IRB.conf[:VERBOSE] = false
+      irb = IRB::Irb.new(workspace, input)
+      IRB.conf[:MAIN_CONTEXT] = irb.context
+      irb.context.return_format = "=> %s\n"
+      out, err = capture_output do
+        irb.eval_input
+      end
+      assert_empty err
+      assert_match(%r[/irb\.rb], out)
+    end
+
     def test_show_source_end_finder
       pend if RUBY_ENGINE == 'truffleruby'
       eval(code = <<-EOS, binding, __FILE__, __LINE__ + 1)
@@ -610,5 +628,45 @@ module TestIRB
       assert_empty err
       assert_match(/^From: .+ @ line \d+ :\n/, out)
     end
+
+    def test_whereami_alias
+      input = TestInputMethod.new([
+        "@\n",
+      ])
+      IRB.init_config(nil)
+      IRB.conf[:COMMAND_ALIASES] = { :'@' => :whereami }
+      workspace = IRB::WorkSpace.new(Object.new)
+      irb = IRB::Irb.new(workspace, input)
+      IRB.conf[:MAIN_CONTEXT] = irb.context
+      out, err = capture_output do
+        irb.eval_input
+      end
+      assert_empty err
+      assert_match(/^From: .+ @ line \d+ :\n/, out)
+    end
+
+    def test_vars_with_aliases
+      input = TestInputMethod.new([
+        "@foo\n",
+        "$bar\n",
+      ])
+      IRB.init_config(nil)
+      IRB.conf[:COMMAND_ALIASES] = {
+        :'@' => :whereami,
+        :'$' => :show_source,
+      }
+      main = Object.new
+      main.instance_variable_set(:@foo, "foo")
+      $bar = "bar"
+      workspace = IRB::WorkSpace.new(main)
+      irb = IRB::Irb.new(workspace, input)
+      IRB.conf[:MAIN_CONTEXT] = irb.context
+      out, err = capture_output do
+        irb.eval_input
+      end
+      assert_empty err
+      assert_match(/"foo"/, out)
+      assert_match(/"bar"/, out)
+    end
   end
 end