From 687e6773bc33825631ebcf6cd05d671a6a7316b2 Mon Sep 17 00:00:00 2001 From: tompng Date: Mon, 18 May 2026 03:12:06 +0900 Subject: [PATCH] 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. --- lib/irb.rb | 2 - lib/irb/color.rb | 21 +++++++-- lib/irb/completion.rb | 24 ++++++++-- lib/irb/context.rb | 2 +- lib/irb/ext/change-ws.rb | 1 - lib/irb/ext/workspaces.rb | 1 - lib/irb/helper_method.rb | 83 +++++++++++++++++++++++++++++++-- lib/irb/input-method.rb | 19 ++++++++ lib/irb/workspace.rb | 24 +--------- test/irb/test_command.rb | 9 ---- test/irb/test_completion.rb | 23 +++++++++ test/irb/test_helper_method.rb | 59 +++++++++++++++++++++++ test/irb/test_type_completor.rb | 18 +++++++ 13 files changed, 237 insertions(+), 49 deletions(-) diff --git a/lib/irb.rb b/lib/irb.rb index af65c8a13..7f6bf5c38 100644 --- a/lib/irb.rb +++ b/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: diff --git a/lib/irb/color.rb b/lib/irb/color.rb index 3e9b59532..7046b464e 100644 --- a/lib/irb/color.rb +++ b/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 diff --git a/lib/irb/completion.rb b/lib/irb/completion.rb index 3dc2fa22a..2a73cfc74 100644 --- a/lib/irb/completion.rb +++ b/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:) diff --git a/lib/irb/context.rb b/lib/irb/context.rb index 284946be0..959ec58f2 100644 --- a/lib/irb/context.rb +++ b/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) diff --git a/lib/irb/ext/change-ws.rb b/lib/irb/ext/change-ws.rb index 60e8afe31..022336dfd 100644 --- a/lib/irb/ext/change-ws.rb +++ b/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 diff --git a/lib/irb/ext/workspaces.rb b/lib/irb/ext/workspaces.rb index da09faa83..00f3319da 100644 --- a/lib/irb/ext/workspaces.rb +++ b/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 diff --git a/lib/irb/helper_method.rb b/lib/irb/helper_method.rb index f1f6fff91..0df5ee326 100644 --- a/lib/irb/helper_method.rb +++ b/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,85 @@ 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 }`. + 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) + # Avoid `{x:conf}` being transformed to `{x:::IRB::HelperMethod::Container.conf}` which is a syntax error + injected << ' ' if injected.end_with?(':') + 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) diff --git a/lib/irb/input-method.rb b/lib/irb/input-method.rb index 32bdce14f..6580388c4 100644 --- a/lib/irb/input-method.rb +++ b/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 diff --git a/lib/irb/workspace.rb b/lib/irb/workspace.rb index 9fef8f86a..43075410b 100644 --- a/lib/irb/workspace.rb +++ b/lib/irb/workspace.rb @@ -96,17 +96,9 @@ def initialize(*main) # IRB.conf[:__MAIN__] 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< #{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", diff --git a/test/irb/test_completion.rb b/test/irb/test_completion.rb index 8959604cb..3de9e9941 100644 --- a/test/irb/test_completion.rb +++ b/test/irb/test_completion.rb @@ -43,6 +43,29 @@ def test_command_document_target end end + class HelperMethodCompletionTest < CompletionTest + def test_helper_method_completion + completor = IRB::RegexpCompletor.new + # Assuming `conf` is a helper method, it should be included in the completion candidates + assert_include(completor.completion_candidates('', 'co', '', bind: binding), 'conf') + assert_include(completor.completion_candidates('p(', 'co', '', bind: binding), 'conf') + assert_not_include(completor.completion_candidates('def f(', 'co', '', bind: binding), 'conf') + end + + def test_helper_method_document_target + completor = IRB::RegexpCompletor.new + result = completor.doc_namespace('tap do ', 'conf', '', bind: binding) + assert_instance_of(IRB::HelperMethodDocument, result) + assert_equal('conf', result.name) + + result = completor.doc_namespace('tap do ', 'conf', '', bind: eval('conf = 1; binding')) + refute_instance_of(IRB::HelperMethodDocument, result) + + result = completor.doc_namespace('tap do |conf| ', 'conf', '', bind: binding) + refute_instance_of(IRB::HelperMethodDocument, result) + end + end + class MethodCompletionTest < CompletionTest def test_complete_string assert_include(completion_candidates("'foo'.up", binding), "'foo'.upcase") diff --git a/test/irb/test_helper_method.rb b/test/irb/test_helper_method.rb index 4b61397b7..164dbd3ec 100644 --- a/test/irb/test_helper_method.rb +++ b/test/irb/test_helper_method.rb @@ -26,6 +26,65 @@ def test_conf_returns_the_context_object assert_empty err assert_include out, "=> \"irb\"" end + + def test_conf_variations + out, err = execute_lines('p "1:#{conf.ap_name}"; p "2:#{self.then { conf().ap_name }}"; p(x:conf.ap_name+"3")') + + assert_empty err + assert_include out, '"1:irb"' + assert_include out, '"2:irb"' + assert_include out, '"irb3"' + end + + def test_conf_code_injection + assert_equal '::IRB::HelperMethod::Container.conf.ap_name', IRB::HelperMethod.inject_helper_methods('conf.ap_name', local_variables: []) + assert_equal 'conf.ap_name', IRB::HelperMethod.inject_helper_methods('conf.ap_name', local_variables: [:conf]) + assert_equal 'a /conf#{::IRB::HelperMethod::Container.conf}/', IRB::HelperMethod.inject_helper_methods('a /conf#{conf}/', local_variables: []) + assert_equal 'a /::IRB::HelperMethod::Container.conf#{conf}/', IRB::HelperMethod.inject_helper_methods('a /conf#{conf}/', local_variables: [:a]) + assert_equal( + '::IRB::HelperMethod::Container.conf.ap_name; conf = 1; conf.ap_name; class A; ::IRB::HelperMethod::Container.conf.ap_name; end', + IRB::HelperMethod.inject_helper_methods('conf.ap_name; conf = 1; conf.ap_name; class A; conf.ap_name; end') + ) + end + + def test_conf_completion + assert_include IRB::HelperMethod.completions('loop do |_conf| ', 'co', local_variables: []), 'conf' + assert_include IRB::HelperMethod.completions('def f(x=', 'co', local_variables: []), 'conf' + assert_not_include IRB::HelperMethod.completions('def f(', 'co', local_variables: []), 'conf' + assert_include IRB::HelperMethod.completions("a /1#/i;'\n", 'co', local_variables: [:a]), 'conf' + assert_not_include IRB::HelperMethod.completions("a /1#/i;'\n", 'co', local_variables: []), 'conf' + end + end + + class ColorizationTest < HelperMethodTestCase + def test_colorize_helper_method + # Without helper_methods: used in inspect result + assert_equal( + "\e[36mconf\e[0m[]; \e[36mconf\e[0m(); +\e[36mconf\e[0m", + IRB::Color.colorize_code('conf[]; conf(); +conf', colorable: true) + ) + + # With helper_methods: used in syntax highlighting + assert_equal( + "\e[1mconf\e[0m[]; \e[1mconf\e[0m(); +\e[1mconf\e[0m", + IRB::Color.colorize_code('conf[]; conf(); +conf', colorable: true, helper_methods: true) + ) + + # If receiver exists, it's not a helper method + assert_not_include(IRB::Color.colorize_code('tap{self.conf}', colorable: true, helper_methods: true), "\e[1mconf\e[0m") + # ImplicitNode is not supported + assert_not_include(IRB::Color.colorize_code('p(conf:)', colorable: true, helper_methods: true), "\e[1mconf\e[0m") + end + + def test_colorize_legacy_command_bundle_helper_method + IRB::ExtendCommandBundle.define_method(:my_helper) {} + assert_equal( + "\e[1mmy_helper\e[0m[]; \e[1mmy_helper\e[0m(); +\e[1mmy_helper\e[0m", + IRB::Color.colorize_code('my_helper[]; my_helper(); +my_helper', colorable: true, helper_methods: true) + ) + ensure + IRB::ExtendCommandBundle.remove_method(:my_helper) + end end end diff --git a/test/irb/test_type_completor.rb b/test/irb/test_type_completor.rb index 910c97c57..f007ce90a 100644 --- a/test/irb/test_type_completor.rb +++ b/test/irb/test_type_completor.rb @@ -96,6 +96,24 @@ def test_command_document_target refute_instance_of(IRB::CommandDocument, result) end + def test_helper_method_completion + assert_include(@completor.completion_candidates('', 'co', '', bind: binding), 'conf') + assert_include(@completor.completion_candidates('p(', 'co', '', bind: binding), 'conf') + assert_not_include(@completor.completion_candidates('def f(', 'co', '', bind: binding), 'conf') + end + + def test_helper_method_document_target + result = @completor.doc_namespace('tap do ', 'conf', '', bind: binding) + assert_instance_of(IRB::HelperMethodDocument, result) + assert_equal('conf', result.name) + + result = @completor.doc_namespace('tap do ', 'conf', '', bind: eval('conf = 1; binding')) + refute_instance_of(IRB::HelperMethodDocument, result) + + result = @completor.doc_namespace('tap do |conf| ', 'conf', '', bind: binding) + refute_instance_of(IRB::HelperMethodDocument, result) + end + def test_type_completor_handles_encoding_errors_gracefully invalid_method_name = "b\xff".dup.force_encoding(Encoding::ASCII_8BIT)