[ruby/irb] Add edit command (https://github.com/ruby/irb/pull/453)

* Add edit command

* Make find_source a public singleton method

* Add document for the edit command

* Make find_end private

* Remove duplicated private

https://github.com/ruby/irb/commit/4321674aa7
Co-authored-by: Takashi Kokubun <takashikkbn@gmail.com>

5 files changed, 197 insertions, 44 deletions

diff --git a/README.md b/README.md
index 5b1f21abc9..9b4249beb2 100644
--- a/README.md
+++ b/README.md
@@ -64,6 +64,13 @@ The following commands are available on IRB.
   * Change the current workspace to an object.
 * `bindings`, `workspaces`
   * Show workspaces.
+* `edit`
+  * Open a file with the editor command defined with `ENV["EDITOR"]`
+    * `edit` - opens the file the current context belongs to (if applicable)
+    * `edit foo.rb` - opens `foo.rb`
+    * `edit Foo` - opens the location of `Foo`
+    * `edit Foo.bar` - opens the location of `Foo.bar`
+    * `edit Foo#bar` - opens the location of `Foo#bar`
 * `pushb`, `pushws`
   * Push an object to the workspace stack.
 * `popb`, `popws`
diff --git a/lib/irb/cmd/edit.rb b/lib/irb/cmd/edit.rb
new file mode 100644
index 0000000000..8d3fab3273
--- /dev/null
+++ b/lib/irb/cmd/edit.rb
@@ -0,0 +1,65 @@
+require 'shellwords'
+require_relative "nop"
+
+module IRB
+  # :stopdoc:
+
+  module ExtendCommand
+    class Edit < Nop
+      class << self
+        def transform_args(args)
+          # Return a string literal as is for backward compatibility
+          if args.nil? || args.empty? || string_literal?(args)
+            args
+          else # Otherwise, consider the input as a String for convenience
+            args.strip.dump
+          end
+        end
+
+        private
+
+        def string_literal?(args)
+          sexp = Ripper.sexp(args)
+          sexp && sexp.size == 2 && sexp.last&.first&.first == :string_literal
+        end
+      end
+
+      def execute(*args)
+        path = args.first
+
+        if path.nil? && (irb_path = @irb_context.irb_path)
+          path = irb_path
+        end
+
+        if !File.exist?(path)
+          require_relative "show_source"
+
+          source =
+            begin
+              ShowSource.find_source(path, @irb_context)
+            rescue NameError
+              # if user enters a path that doesn't exist, it'll cause NameError when passed here because find_source would try to evaluate it as well
+              # in this case, we should just ignore the error
+            end
+
+          if source && File.exist?(source.file)
+            path = source.file
+          else
+            puts "Can not find file: #{path}"
+            return
+          end
+        end
+
+        if editor = ENV['EDITOR']
+          puts "command: '#{editor}'"
+          puts "   path: #{path}"
+          system(*Shellwords.split(editor), path)
+        else
+          puts "Can not find editor setting: ENV['EDITOR']"
+        end
+      end
+    end
+  end
+
+  # :startdoc:
+end
diff --git a/lib/irb/cmd/show_source.rb b/lib/irb/cmd/show_source.rb
index 1fcff3e897..03c21b78c7 100644
--- a/lib/irb/cmd/show_source.rb
+++ b/lib/irb/cmd/show_source.rb
@@ -19,8 +19,50 @@ module IRB
           end
         end
 
+        def find_source(str, irb_context)
+          case str
+          when /\A[A-Z]\w*(::[A-Z]\w*)*\z/ # Const::Name
+            eval(str, irb_context.workspace.binding) # trigger autoload
+            base = irb_context.workspace.binding.receiver.yield_self { |r| r.is_a?(Module) ? r : Object }
+            file, line = base.const_source_location(str) if base.respond_to?(:const_source_location) # Ruby 2.7+
+          when /\A(?<owner>[A-Z]\w*(::[A-Z]\w*)*)#(?<method>[^ :.]+)\z/ # Class#method
+            owner = eval(Regexp.last_match[:owner], irb_context.workspace.binding)
+            method = Regexp.last_match[:method]
+            if owner.respond_to?(:instance_method) && owner.instance_methods.include?(method.to_sym)
+              file, line = owner.instance_method(method).source_location
+            end
+          when /\A((?<receiver>.+)(\.|::))?(?<method>[^ :.]+)\z/ # method, receiver.method, receiver::method
+            receiver = eval(Regexp.last_match[:receiver] || 'self', irb_context.workspace.binding)
+            method = Regexp.last_match[:method]
+            file, line = receiver.method(method).source_location if receiver.respond_to?(method)
+          end
+          if file && line
+            Source.new(file: file, first_line: line, last_line: find_end(file, line))
+          end
+        end
+
         private
 
+        def find_end(file, first_line)
+          return first_line unless File.exist?(file)
+          lex = RubyLex.new
+          lines = File.read(file).lines[(first_line - 1)..-1]
+          tokens = RubyLex.ripper_lex_without_warning(lines.join)
+          prev_tokens = []
+
+          # chunk with line number
+          tokens.chunk { |tok| tok.pos[0] }.each do |lnum, chunk|
+            code = lines[0..lnum].join
+            prev_tokens.concat chunk
+            continue = lex.process_continue(prev_tokens)
+            code_block_open = lex.check_code_block(code, prev_tokens)
+            if !continue && !code_block_open
+              return first_line + lnum
+            end
+          end
+          first_line
+        end
+
         def string_literal?(args)
           sexp = Ripper.sexp(args)
           sexp && sexp.size == 2 && sexp.last&.first&.first == :string_literal
@@ -32,7 +74,8 @@ module IRB
           puts "Error: Expected a string but got #{str.inspect}"
           return
         end
-        source = find_source(str)
+
+        source = self.class.find_source(str, @irb_context)
         if source && File.exist?(source.file)
           show_source(source)
         else
@@ -53,48 +96,6 @@ module IRB
         puts
       end
 
-      def find_source(str)
-        case str
-        when /\A[A-Z]\w*(::[A-Z]\w*)*\z/ # Const::Name
-          eval(str, irb_context.workspace.binding) # trigger autoload
-          base = irb_context.workspace.binding.receiver.yield_self { |r| r.is_a?(Module) ? r : Object }
-          file, line = base.const_source_location(str) if base.respond_to?(:const_source_location) # Ruby 2.7+
-        when /\A(?<owner>[A-Z]\w*(::[A-Z]\w*)*)#(?<method>[^ :.]+)\z/ # Class#method
-          owner = eval(Regexp.last_match[:owner], irb_context.workspace.binding)
-          method = Regexp.last_match[:method]
-          if owner.respond_to?(:instance_method) && owner.instance_methods.include?(method.to_sym)
-            file, line = owner.instance_method(method).source_location
-          end
-        when /\A((?<receiver>.+)(\.|::))?(?<method>[^ :.]+)\z/ # method, receiver.method, receiver::method
-          receiver = eval(Regexp.last_match[:receiver] || 'self', irb_context.workspace.binding)
-          method = Regexp.last_match[:method]
-          file, line = receiver.method(method).source_location if receiver.respond_to?(method)
-        end
-        if file && line
-          Source.new(file: file, first_line: line, last_line: find_end(file, line))
-        end
-      end
-
-      def find_end(file, first_line)
-        return first_line unless File.exist?(file)
-        lex = RubyLex.new
-        lines = File.read(file).lines[(first_line - 1)..-1]
-        tokens = RubyLex.ripper_lex_without_warning(lines.join)
-        prev_tokens = []
-
-        # chunk with line number
-        tokens.chunk { |tok| tok.pos[0] }.each do |lnum, chunk|
-          code = lines[0..lnum].join
-          prev_tokens.concat chunk
-          continue = lex.process_continue(prev_tokens)
-          code_block_open = lex.check_code_block(code, prev_tokens)
-          if !continue && !code_block_open
-            return first_line + lnum
-          end
-        end
-        first_line
-      end
-
       def bold(str)
         Color.colorize(str, [:BOLD])
       end
diff --git a/lib/irb/extend-command.rb b/lib/irb/extend-command.rb
index 94fd9c8bb4..7e120cf510 100644
--- a/lib/irb/extend-command.rb
+++ b/lib/irb/extend-command.rb
@@ -121,6 +121,10 @@ module IRB # :nodoc:
         [:debug, NO_OVERRIDE],
       ],
       [
+        :irb_edit, :Edit, "cmd/edit",
+        [:edit, NO_OVERRIDE],
+      ],
+      [
         :irb_help, :Help, "cmd/help",
         [:help, NO_OVERRIDE],
       ],
diff --git a/test/irb/test_cmd.rb b/test/irb/test_cmd.rb
index bcfb1d0b86..44f348a724 100644
--- a/test/irb/test_cmd.rb
+++ b/test/irb/test_cmd.rb
@@ -565,9 +565,84 @@ module TestIRB
       $bar = nil
     end
 
+    class EditTest < ExtendCommandTest
+      def setup
+        @original_editor = ENV["EDITOR"]
+        # noop the command so nothing gets executed
+        ENV["EDITOR"] = ": code"
+      end
+
+      def teardown
+        ENV["EDITOR"] = @original_editor
+      end
+
+      def test_edit_without_arg
+        out, err = execute_lines(
+          "edit",
+          irb_path: __FILE__
+        )
+
+        assert_empty err
+        assert_match("path: #{__FILE__}", out)
+        assert_match("command: ': code'", out)
+      end
+
+      def test_edit_with_path
+        out, err = execute_lines(
+          "edit #{__FILE__}"
+        )
+
+        assert_empty err
+        assert_match("path: #{__FILE__}", out)
+        assert_match("command: ': code'", out)
+      end
+
+      def test_edit_with_non_existing_path
+        out, err = execute_lines(
+          "edit foo.rb"
+        )
+
+        assert_empty err
+        assert_match /Can not find file: foo\.rb/, out
+      end
+
+      def test_edit_with_constant
+        # const_source_location is supported after Ruby 2.7
+        omit if Gem::Version.new(RUBY_VERSION) < Gem::Version.new('2.7.0') || RUBY_ENGINE == 'truffleruby'
+
+        out, err = execute_lines(
+          "edit IRB::Irb"
+        )
+
+        assert_empty err
+        assert_match(/path: .*\/lib\/irb\.rb/, out)
+        assert_match("command: ': code'", out)
+      end
+
+      def test_edit_with_class_method
+        out, err = execute_lines(
+          "edit IRB.start"
+        )
+
+        assert_empty err
+        assert_match(/path: .*\/lib\/irb\.rb/, out)
+        assert_match("command: ': code'", out)
+      end
+
+      def test_edit_with_instance_method
+        out, err = execute_lines(
+          "edit IRB::Irb#run"
+        )
+
+        assert_empty err
+        assert_match(/path: .*\/lib\/irb\.rb/, out)
+        assert_match("command: ': code'", out)
+      end
+    end
+
     private
 
-    def execute_lines(*lines, conf: {}, main: self)
+    def execute_lines(*lines, conf: {}, main: self, irb_path: nil)
       IRB.init_config(nil)
       IRB.conf[:VERBOSE] = false
       IRB.conf[:PROMPT_MODE] = :SIMPLE
@@ -575,6 +650,7 @@ module TestIRB
       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
       capture_output do
         irb.eval_input