Allow non-identifier aliases like Pry's @ and $ (#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

lib/irb.rb

@@ -426,6 +426,9 @@ class 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

lib/irb/context.rb

@@ -149,6 +149,8 @@ def initialize(irb, workspace = nil, input_method = nil)
       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 @@ def main
     # 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 @@ def evaluate(line, line_no, exception: nil) # :nodoc:
         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 @@ def inspect # :nodoc:
     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

lib/irb/init.rb

@@ -158,6 +158,8 @@ def IRB.init_config(ap_path)
     @CONF[:LC_MESSAGES] = Locale.new
 
     @CONF[:AT_EXIT] = []
+
+    @CONF[:COMMAND_ALIASES] = {}
   end
 
   def IRB.set_measure_callback(type = nil, arg = nil, &block)

lib/irb/ruby-lex.rb

@@ -65,6 +65,12 @@ def set_input(io, p = nil, context:, &block)
             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

test/irb/test_cmd.rb

@@ -570,6 +570,24 @@ def test_show_source
       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 @@ def test_whereami
       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