Helper methods without object polluting

Instead of defining helper method to current workspace object, modify the input code to directly call helper method defined in a container object.
`p conf.ap_name` will be modified to `p ::IRB::HelperMethod::Container.conf.ap_name`.
Also implements highlighting, completion, and document dialog for helper methods.

Author: tompng

lib/irb.rb

@@ -105,7 +105,6 @@ def initialize(workspace = nil, input_method = nil, from_binding: false)
       @from_binding = from_binding
       @prompt_part_cache = nil
       @context = Context.new(self, workspace, input_method)
-      @context.workspace.load_helper_methods_to_main
       @signal_status = :IN_IRB
       @scanner = RubyLex.new
       @line_no = 1
@@ -126,7 +125,6 @@ def debug_break
     def debug_readline(binding)
       workspace = IRB::WorkSpace.new(binding)
       context.replace_workspace(workspace)
-      context.workspace.load_helper_methods_to_main
       @line_no += 1
 
       # When users run:

lib/irb/color.rb

@@ -18,6 +18,11 @@ module Color
     CYAN      = 36
     WHITE     = 37
 
+    PRIORITY_ERROR = 0
+    PRIORITY_HELPER_METHOD = 1
+    PRIORITY_PRIOR_TOKEN = 2
+    PRIORITY_NORMAL_TOKEN = 3
+
     # Following pry's colors where possible
     TOKEN_SEQS = {
       KEYWORD_NIL:        [CYAN, BOLD],
@@ -106,6 +111,8 @@ module Color
       method_name:        [CYAN, BOLD],
       message_name:       [CYAN],
       symbol:             [YELLOW],
+      # helper method
+      helper_method:      [BOLD],
       # special colorization
       error:              [RED, REVERSE],
     }.transform_values do |styles|
@@ -162,7 +169,7 @@ def colorize(text, seq, colorable: colorable?)
       # If `complete` is false (code is incomplete), this does not warn compile_error.
       # This option is needed to avoid warning a user when the compile_error is happening
       # because the input is not wrong but just incomplete.
-      def colorize_code(code, complete: true, ignore_error: false, colorable: colorable?, local_variables: [])
+      def colorize_code(code, complete: true, ignore_error: false, colorable: colorable?, local_variables: [], helper_methods: false)
         return code unless colorable
 
         result = Prism.parse_lex(code, scopes: [local_variables])
@@ -180,9 +187,13 @@ def colorize_code(code, complete: true, ignore_error: false, colorable: colorabl
         visitor = ColorizeVisitor.new
         prism_node.accept(visitor)
 
-        error_tokens = errors.map { |e| [e.location.start_line, e.location.start_column, 0, e.location.end_line, e.location.end_column, :error, e.location.slice] }
+        helper_method_locations = helper_methods ? IRB::HelperMethod.extract_helper_method_locations(prism_node) : []
+        helper_method_tokens = helper_method_locations.map { |loc| [loc.start_line, loc.start_column, PRIORITY_HELPER_METHOD, loc.end_line, loc.end_column, :helper_method, loc.slice] }
+
+        error_tokens = errors.map { |e| [e.location.start_line, e.location.start_column, PRIORITY_ERROR, e.location.end_line, e.location.end_column, :error, e.location.slice] }
         error_tokens.reject! { |t| t.last.match?(/\A\s*\z/) }
-        tokens = prism_tokens.map { |t,| [t.location.start_line, t.location.start_column, 2, t.location.end_line, t.location.end_column, t.type, t.value] }
+
+        tokens = prism_tokens.map { |t,| [t.location.start_line, t.location.start_column, PRIORITY_NORMAL_TOKEN, t.location.end_line, t.location.end_column, t.type, t.value] }
         tokens.pop if tokens.last&.[](5) == :EOF
 
         colored = +''
@@ -201,7 +212,7 @@ def colorize_code(code, complete: true, ignore_error: false, colorable: colorabl
           end
         }
 
-        (visitor.tokens + tokens + error_tokens).sort.each do |start_line, start_column, _priority, end_line, end_column, type, value|
+        (visitor.tokens + tokens + error_tokens + helper_method_tokens).sort.each do |start_line, start_column, _priority, end_line, end_column, type, value|
           next if start_line - 1 < line_index || (start_line - 1 == line_index && start_column < col)
 
           flush.call(start_line - 1, start_column)
@@ -235,7 +246,7 @@ def initialize
 
         def dispatch(location, type)
           if location
-            @tokens << [location.start_line, location.start_column, 1, location.end_line, location.end_column, type, location.slice]
+            @tokens << [location.start_line, location.start_column, PRIORITY_PRIOR_TOKEN, location.end_line, location.end_column, type, location.slice]
           end
         end
 

lib/irb/completion.rb

@@ -19,6 +19,9 @@ def initialize(name)
   class CommandDocument < DocumentTarget # :nodoc:
   end
 
+  class HelperMethodDocument < DocumentTarget # :nodoc:
+  end
+
   # Represents a method/class documentation target. May hold multiple names
   # when the receiver is ambiguous (e.g. `{}.any?` could be Hash#any? or Proc#any?).
   # The dialog popup uses only the first name; the full-screen display renders all.
@@ -105,6 +108,12 @@ def command_document_target(preposing, matched)
       end
     end
 
+    def helper_method_document_target(preposing, matched, local_variables:)
+      if IRB::HelperMethod.completions(preposing, matched, local_variables: local_variables).include?(matched)
+        HelperMethodDocument.new(matched)
+      end
+    end
+
     def retrieve_files_to_require_relative_from_current_dir
       @files_from_current_dir ||= Dir.glob("**/*.{rb,#{RbConfig::CONFIG['DLEXT']}}", base: '.').map { |path|
         path.sub(/\.(rb|#{RbConfig::CONFIG['DLEXT']})\z/, '')
@@ -143,11 +152,14 @@ def completion_candidates(preposing, target, _postposing, bind:)
         # If the string cannot be converted, we just ignore it
         nil
       end
-      commands | encoded_candidates
+      helper_methods = IRB::HelperMethod.completions(preposing, target, local_variables: bind.local_variables)
+      commands | helper_methods | encoded_candidates
     end
 
     def doc_namespace(preposing, matched, _postposing, bind:)
-      command_document_target(preposing, matched) || begin
+      command_document_target(preposing, matched) ||
+      helper_method_document_target(preposing, matched, local_variables: bind.local_variables) ||
+      begin
         result = ReplTypeCompletor.analyze(preposing + matched, binding: bind, filename: @context.irb_path)
         result&.doc_namespace('')
       end
@@ -229,11 +241,15 @@ def completion_candidates(preposing, target, postposing, bind:)
         # If the string cannot be converted, we just ignore it
         nil
       end
-      commands | completion_data
+
+      helper_methods = IRB::HelperMethod.completions(preposing, target, local_variables: bind.local_variables)
+      commands | helper_methods | completion_data
     end
 
     def doc_namespace(preposing, matched, _postposing, bind:)
-      command_document_target(preposing, matched) || retrieve_completion_data(matched, bind: bind, doc_namespace: true)
+      command_document_target(preposing, matched) ||
+      helper_method_document_target(preposing, matched, local_variables: bind.local_variables) ||
+      retrieve_completion_data(matched, bind: bind, doc_namespace: true)
     end
 
     def retrieve_completion_data(input, bind:, doc_namespace:)

lib/irb/context.rb

@@ -652,7 +652,7 @@ def colorize_input(input, complete:)
           arg = IRB::Color.colorize_code(arg, complete: complete, local_variables: lvars) if arg
           "#{IRB::Color.colorize(name, [:BOLD])}\e[m#{sep}#{arg}"
         else
-          IRB::Color.colorize_code(input, complete: complete, local_variables: lvars)
+          IRB::Color.colorize_code(input, complete: complete, local_variables: lvars, helper_methods: true)
         end
       else
         Reline::Unicode.escape_for_print(input)

lib/irb/ext/change-ws.rb

@@ -31,7 +31,6 @@ def change_workspace(*_main)
 
       workspace = WorkSpace.new(_main[0])
       replace_workspace(workspace)
-      workspace.load_helper_methods_to_main
     end
   end
 end

lib/irb/ext/workspaces.rb

@@ -21,7 +21,6 @@ def push_workspace(*_main)
       else
         new_workspace = WorkSpace.new(workspace.binding, _main[0])
         @workspace_stack.push new_workspace
-        new_workspace.load_helper_methods_to_main
       end
     end
 

lib/irb/helper_method.rb

@@ -1,4 +1,5 @@
 require_relative "helper_method/base"
+require "prism"
 
 module IRB
   module HelperMethod
@@ -9,9 +10,8 @@ class << self
 
       def register(name, helper_class)
         @helper_methods[name] = helper_class
-
-        if defined?(HelpersContainer)
-          HelpersContainer.install_helper_methods
+        Container.define_singleton_method name do |*args, **opts, &block|
+          helper_class.instance.execute(*args, **opts, &block)
         end
       end
 
@@ -20,8 +20,84 @@ def all_helper_methods_info
           { display_name: name, description: helper_class.description }
         end
       end
+
+      # Injects helper method calls with the corresponding container method calls.
+      # For example, `tap { p conf.ap_name }` will be transformed to `tap { p IRB::HelperMethod::Container.conf.ap_name }`.
+      # Ideally, we want to use `::IRB::HelperMethod::Container`, but `{ key:conf }` can't be parsed if we prefix with `::`.
+      def inject_helper_methods(code, local_variables: [])
+        parse_result = Prism.parse(code, scopes: [local_variables])
+        return code unless parse_result.success?
+
+        locations = extract_helper_method_locations(parse_result.value)
+        return code if locations.empty?
+
+        injected = +''
+        offset = 0
+        locations.each do |loc|
+          injected << code.byteslice(offset...loc.start_offset)
+          injected << "IRB::HelperMethod::Container.#{loc.slice}"
+          offset = loc.end_offset
+        end
+        injected << code.byteslice(offset..)
+        injected
+      end
+
+      def completions(preposing, target, local_variables:)
+        helper_method_names = @helper_methods.keys.map(&:to_s)
+        candidates = helper_method_names.select {|name| name.start_with?(target) }
+        return [] if candidates.empty?
+
+        target_message = nil
+        end_offset = preposing.bytesize + target.bytesize
+        visitor = MethodCallVisitor.new do |call_node|
+          target_message = call_node.message if call_node.message_loc.end_offset == end_offset
+        end
+        Prism.parse(preposing + target, scopes: [local_variables]).value.accept(visitor)
+        return [] unless target_message
+
+        candidates
+      end
+
+      def extract_helper_method_locations(node)
+        helper_method_names = @helper_methods.keys.map(&:to_s)
+
+        # Legacy helper methods defined in ExtendCommandBundle should also be considered as helper methods
+        helper_method_names += IRB::ExtendCommandBundle.instance_methods.map(&:to_s)
+
+        helper_method_locations = []
+        visitor = MethodCallVisitor.new do |call_node|
+          if helper_method_names.include?(call_node.message)
+            helper_method_locations << call_node.message_loc
+          end
+        end
+        visitor.visit(node)
+        helper_method_locations.sort_by(&:start_offset)
+      end
+    end
+
+    # Traverse and finds CallNode without receiver which may be helper method calls.
+    class MethodCallVisitor < Prism::Visitor # :nodoc:
+      def initialize(&block)
+        @callback = block
+      end
+
+      def visit_call_node(node)
+        super
+        @callback.call node if node.receiver.nil?
+      end
+
+      def visit_implicit_node(node)
+        # We can't modify `{ conf: }` to `{ IRB::HelperMethod::Container.conf: }`
+        # so it can't be a helper method call
+      end
     end
 
+    Container = Object.new
+
+    # Enable legacy helper method registration for backward compatibility
+    require_relative "default_commands"
+    Container.extend IRB::ExtendCommandBundle
+
     # Default helper_methods
     require_relative "helper_method/conf"
     register(:conf, HelperMethod::Conf)

lib/irb/input-method.rb

@@ -374,6 +374,8 @@ def show_doc_dialog_proc
         contents = case target
         when CommandDocument
           input_method.command_doc_dialog_contents(target.name, width)
+        when HelperMethodDocument
+          input_method.helper_method_doc_dialog_contents(target.name, width)
         when MethodDocument
           input_method.rdoc_dialog_contents(target.name, width)
         else
@@ -396,6 +398,16 @@ def command_doc_dialog_contents(command_name, width)
       [PRESS_ALT_D_TO_READ_FULL_DOC, ""] + command_class.doc_dialog_content(command_name, width)
     end
 
+    def helper_method_doc_dialog_contents(helper_method_name, width)
+      helper_method_class = IRB::HelperMethod.helper_methods[helper_method_name.to_sym]
+      return unless helper_method_class
+      [
+        PRESS_ALT_D_TO_READ_FULL_DOC, "",
+        Color.colorize(helper_method_name, [:BOLD, :BLUE]) + Color.colorize(" (helper method)", [:CYAN]), "",
+        helper_method_class.description
+      ]
+    end
+
     def easter_egg_dialog_contents
       type = STDOUT.external_encoding == Encoding::UTF_8 ? :unicode : :ascii
       lines = IRB.send(:easter_egg_logo, type).split("\n")
@@ -474,6 +486,13 @@ def display_document(matched)
             io.puts content
           end
         end
+      when HelperMethodDocument
+        helper_method_class = IRB::HelperMethod.helper_methods[target.name.to_sym]
+        if helper_method_class
+          Pager.page(retain_content: true) do |io|
+            io.puts helper_method_class.description
+          end
+        end
       when MethodDocument
         driver = rdoc_ri_driver
         return unless driver

lib/irb/workspace.rb

@@ -96,17 +96,9 @@ def initialize(*main)
     # <code>IRB.conf[:__MAIN__]</code>
     attr_reader :main
 
-    def load_helper_methods_to_main
-      # Do not load helper methods to frozen objects and BasicObject
-      return unless Object === @main && !@main.frozen?
-
-      ancestors = class<<main;ancestors;end
-      main.extend ExtendCommandBundle if !ancestors.include?(ExtendCommandBundle)
-      main.extend HelpersContainer if !ancestors.include?(HelpersContainer)
-    end
-
     # Evaluate the given +statements+ within the  context of this workspace.
     def evaluate(statements, file = __FILE__, line = __LINE__)
+      statements = HelperMethod.inject_helper_methods(statements, local_variables: @binding.local_variables)
       eval(statements, @binding, file, line)
     end
 
@@ -163,18 +155,4 @@ def code_around_binding
       "\nFrom: #{file} @ line #{pos + 1} :\n\n#{body}#{Color.clear}\n"
     end
   end
-
-  module HelpersContainer
-    class << self
-      def install_helper_methods
-        HelperMethod.helper_methods.each do |name, helper_method_class|
-          define_method name do |*args, **opts, &block|
-            helper_method_class.instance.execute(*args, **opts, &block)
-          end unless method_defined?(name)
-        end
-      end
-    end
-
-    install_helper_methods
-  end
 end

test/irb/test_command.rb

@@ -495,15 +495,6 @@ def test_pushws_switches_to_new_workspace_and_pushes_the_current_one_to_the_stac
       assert_match(/=> #{self.class}\n$/, out)
     end
 
-    def test_pushws_extends_the_new_workspace_with_command_bundle
-      out, err = execute_lines(
-        "pushws Object.new",
-        "self.singleton_class.ancestors"
-      )
-      assert_empty err
-      assert_include(out, "IRB::ExtendCommandBundle")
-    end
-
     def test_pushws_prints_workspace_stack_when_no_arg_is_given
       out, err = execute_lines(
         "pushws",