[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

3 files changed, 26 insertions, 0 deletions

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