From 20b439506635f13d1296072a13f0dc6a5f4db542 Mon Sep 17 00:00:00 2001
From: Stan Lo <stan001212@gmail.com>
Date: Mon, 22 Apr 2024 21:54:10 +0100
Subject: [PATCH] Support helper method registration

---
 lib/irb/command/help.rb        |   2 +
 lib/irb/default_commands.rb    |   5 +-
 lib/irb/helper_method.rb       |  29 +++++++++
 lib/irb/helper_method/base.rb  |  12 ++++
 lib/irb/helper_method/conf.rb  |  11 ++++
 lib/irb/workspace.rb           |  20 +++++-
 test/irb/command/test_help.rb  |   9 +++
 test/irb/test_helper_method.rb | 109 +++++++++++++++++++++++++++++++++
 8 files changed, 190 insertions(+), 7 deletions(-)
 create mode 100644 lib/irb/helper_method.rb
 create mode 100644 lib/irb/helper_method/base.rb
 create mode 100644 lib/irb/helper_method/conf.rb
 create mode 100644 test/irb/test_helper_method.rb

diff --git a/lib/irb/command/help.rb b/lib/irb/command/help.rb
index fc44f6080..1ed7a7707 100644
--- a/lib/irb/command/help.rb
+++ b/lib/irb/command/help.rb
@@ -24,7 +24,9 @@ def execute(command_name)
 
       def help_message
         commands_info = IRB::Command.all_commands_info
+        helper_methods_info = IRB::HelperMethod.all_helper_methods_info
         commands_grouped_by_categories = commands_info.group_by { |cmd| cmd[:category] }
+        commands_grouped_by_categories["Helper methods"] = helper_methods_info
 
         user_aliases = irb_context.instance_variable_get(:@user_aliases)
 
diff --git a/lib/irb/default_commands.rb b/lib/irb/default_commands.rb
index 72884318a..2c515674a 100644
--- a/lib/irb/default_commands.rb
+++ b/lib/irb/default_commands.rb
@@ -98,10 +98,7 @@ def load_command(command)
     end
 
     _register_with_aliases(:irb_context, Command::Context,
-      [
-        [:context, NO_OVERRIDE],
-        [:conf, NO_OVERRIDE],
-      ],
+      [:context, NO_OVERRIDE]
     )
 
     _register_with_aliases(:irb_exit, Command::Exit,
diff --git a/lib/irb/helper_method.rb b/lib/irb/helper_method.rb
new file mode 100644
index 000000000..f1f6fff91
--- /dev/null
+++ b/lib/irb/helper_method.rb
@@ -0,0 +1,29 @@
+require_relative "helper_method/base"
+
+module IRB
+  module HelperMethod
+    @helper_methods = {}
+
+    class << self
+      attr_reader :helper_methods
+
+      def register(name, helper_class)
+        @helper_methods[name] = helper_class
+
+        if defined?(HelpersContainer)
+          HelpersContainer.install_helper_methods
+        end
+      end
+
+      def all_helper_methods_info
+        @helper_methods.map do |name, helper_class|
+          { display_name: name, description: helper_class.description }
+        end
+      end
+    end
+
+    # Default helper_methods
+    require_relative "helper_method/conf"
+    register(:conf, HelperMethod::Conf)
+  end
+end
diff --git a/lib/irb/helper_method/base.rb b/lib/irb/helper_method/base.rb
new file mode 100644
index 000000000..dc74b046d
--- /dev/null
+++ b/lib/irb/helper_method/base.rb
@@ -0,0 +1,12 @@
+module IRB
+  module HelperMethod
+    class Base
+      class << self
+        def description(description = nil)
+          @description = description if description
+          @description
+        end
+      end
+    end
+  end
+end
diff --git a/lib/irb/helper_method/conf.rb b/lib/irb/helper_method/conf.rb
new file mode 100644
index 000000000..460f5ab78
--- /dev/null
+++ b/lib/irb/helper_method/conf.rb
@@ -0,0 +1,11 @@
+module IRB
+  module HelperMethod
+    class Conf < Base
+      description "Returns the current context."
+
+      def execute
+        IRB.CurrentContext
+      end
+    end
+  end
+end
diff --git a/lib/irb/workspace.rb b/lib/irb/workspace.rb
index 1490f7b47..dd92d9901 100644
--- a/lib/irb/workspace.rb
+++ b/lib/irb/workspace.rb
@@ -6,6 +6,8 @@
 
 require "delegate"
 
+require_relative "helper_method"
+
 IRB::TOPLEVEL_BINDING = binding
 module IRB # :nodoc:
   class WorkSpace
@@ -109,9 +111,9 @@ def initialize(*main)
     attr_reader :main
 
     def load_helper_methods_to_main
-      if !(class<<main;ancestors;end).include?(ExtendCommandBundle)
-        main.extend ExtendCommandBundle
-      end
+      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.
@@ -172,4 +174,16 @@ def code_around_binding
       "\nFrom: #{file} @ line #{pos + 1} :\n\n#{body}#{Color.clear}\n"
     end
   end
+
+  module HelpersContainer
+    def self.install_helper_methods
+      HelperMethod.helper_methods.each do |name, helper_method_class|
+        define_method name do |*args, **opts, &block|
+          helper_method_class.new.execute(*args, **opts, &block)
+        end unless method_defined?(name)
+      end
+    end
+
+    install_helper_methods
+  end
 end
diff --git a/test/irb/command/test_help.rb b/test/irb/command/test_help.rb
index c82c43a4c..df3753dae 100644
--- a/test/irb/command/test_help.rb
+++ b/test/irb/command/test_help.rb
@@ -62,5 +62,14 @@ def test_help_lists_user_aliases
       assert_match(/\$\s+Alias for `show_source`/, out)
       assert_match(/@\s+Alias for `whereami`/, out)
     end
+
+    def test_help_lists_helper_methods
+      out = run_ruby_file do
+        type "help"
+        type "exit"
+      end
+
+      assert_match(/Helper methods\s+conf\s+Returns the current context/, out)
+    end
   end
 end
diff --git a/test/irb/test_helper_method.rb b/test/irb/test_helper_method.rb
new file mode 100644
index 000000000..5174e5ceb
--- /dev/null
+++ b/test/irb/test_helper_method.rb
@@ -0,0 +1,109 @@
+# frozen_string_literal: true
+require "irb"
+
+require_relative "helper"
+
+module TestIRB
+  class HelperMethodTestCase < TestCase
+    def setup
+      $VERBOSE = nil
+      @verbosity = $VERBOSE
+      save_encodings
+      IRB.instance_variable_get(:@CONF).clear
+    end
+
+    def teardown
+      $VERBOSE = @verbosity
+      restore_encodings
+    end
+
+    def execute_lines(*lines, conf: {}, main: self, irb_path: nil)
+      IRB.init_config(nil)
+      IRB.conf[:VERBOSE] = false
+      IRB.conf[:PROMPT_MODE] = :SIMPLE
+      IRB.conf.merge!(conf)
+      input = TestInputMethod.new(lines)
+      irb = IRB::Irb.new(IRB::WorkSpace.new(main), input)
+      irb.context.return_format = "=> %s\n"
+      irb.context.irb_path = irb_path if irb_path
+      IRB.conf[:MAIN_CONTEXT] = irb.context
+      IRB.conf[:USE_PAGER] = false
+      capture_output do
+        irb.eval_input
+      end
+    end
+  end
+
+  module TestHelperMethod
+    class ConfTest < HelperMethodTestCase
+      def test_conf_returns_the_context_object
+        out, err = execute_lines("conf.ap_name")
+
+        assert_empty err
+        assert_include out, "=> \"irb\""
+      end
+    end
+  end
+
+  class HelperMethodIntegrationTest < IntegrationTestCase
+    def test_arguments_propogation
+      write_ruby <<~RUBY
+        require "irb/helper_method"
+
+        class MyHelper < IRB::HelperMethod::Base
+          description "This is a test helper"
+
+          def execute(
+            required_arg, optional_arg = nil, *splat_arg, required_keyword_arg:,
+            optional_keyword_arg: nil, **double_splat_arg, &block_arg
+          )
+            puts [required_arg, optional_arg, splat_arg, required_keyword_arg, optional_keyword_arg, double_splat_arg, block_arg.call].to_s
+          end
+        end
+
+        IRB::HelperMethod.register(:my_helper, MyHelper)
+
+        binding.irb
+      RUBY
+
+      output = run_ruby_file do
+        type <<~INPUT
+          my_helper(
+            "required", "optional", "splat", required_keyword_arg: "required",
+            optional_keyword_arg: "optional", a: 1, b: 2
+          ) { "block" }
+        INPUT
+        type "exit"
+      end
+
+      assert_include(output, '["required", "optional", ["splat"], "required", "optional", {:a=>1, :b=>2}, "block"]')
+    end
+
+    def test_helper_method_injection_can_happen_after_irb_require
+      write_ruby <<~RUBY
+        require "irb"
+
+        class MyHelper < IRB::HelperMethod::Base
+          description "This is a test helper"
+
+          def execute
+            puts "Hello from MyHelper"
+          end
+        end
+
+        IRB::HelperMethod.register(:my_helper, MyHelper)
+
+        binding.irb
+      RUBY
+
+      output = run_ruby_file do
+        type <<~INPUT
+          my_helper
+        INPUT
+        type "exit"
+      end
+
+      assert_include(output, 'Hello from MyHelper')
+    end
+  end
+end