Preserve Module.nesting on edit-method -p [Fixes #645]

lib/pry/code.rb

@@ -356,6 +356,15 @@ class Pry
       self.class.expression_at(raw, line_number, :consume => consume)
     end
 
+    # Get the (approximate) Module.nesting at the give line number.
+    #
+    # @param [Fixnum]  line_number  line number starting from 1
+    # @param [Module] top_module   the module in which this code exists
+    # @return [Array<Module>]  a list of open modules.
+    def nesting_at(line_number, top_module=Object)
+      Pry::Indent.nesting_at(raw, line_number)
+    end
+
     # Return an unformatted String of the code.
     #
     # @return [String]

lib/pry/default_commands/editing.rb

@@ -213,18 +213,20 @@ class Pry
           lines[0] = definition_line_for_owner(lines[0])
 
           temp_file do |f|
-            f.puts lines.join
+            f.puts lines
             f.flush
             f.close(false)
             invoke_editor(f.path, 0, true)
 
+            source = wrap_for_nesting(wrap_for_owner(File.read(f.path)))
+
             if @method.alias?
               with_method_transaction(original_name, @method.owner) do
-                Pry.new(:input => StringIO.new(File.read(f.path))).rep(@method.owner)
+                Pry.new(:input => StringIO.new(source)).rep(TOPLEVEL_BINDING)
                 Pry.binding_for(@method.owner).eval("alias #{@method.name} #{original_name}")
               end
             else
-              Pry.new(:input => StringIO.new(File.read(f.path))).rep(@method.owner)
+              Pry.new(:input => StringIO.new(source)).rep(TOPLEVEL_BINDING)
             end
           end
         end
@@ -252,7 +254,16 @@ class Pry
             end
           end
 
-          def with_method_transaction(meth_name, target=TOPLEVEL_BINDING)
+          # Run some code ensuring that at the end target#meth_name will not have changed.
+          #
+          # When we're redefining aliased methods we will overwrite the method at the
+          # unaliased name (so that super continues to work). By wrapping that code in a
+          # transation we make that not happen, which means that alias_method_chains, etc.
+          # continue to work.
+          #
+          # @param [String] meth_name  The method name before aliasing
+          # @param [Module] target  The owner of the method
+          def with_method_transaction(meth_name, target)
             target = Pry.binding_for(target)
             temp_name = "__pry_#{meth_name}__"
 
@@ -282,7 +293,6 @@ class Pry
           #
           # @param String   The original definition line. e.g. def self.foo(bar, baz=1)
           # @return String  The new definition line. e.g. def foo(bar, baz=1)
-          #
           def definition_line_for_owner(line)
             if line =~ /^def (?:.*?\.)?#{Regexp.escape(original_name)}(?=[\(\s;]|$)/
               "def #{original_name}#{$'}"
@@ -290,6 +300,39 @@ class Pry
               raise CommandError, "Could not find original `def #{original_name}` line to patch."
             end
           end
+
+          # Update the source code so that when it has the right owner when eval'd.
+          #
+          # This (combined with definition_line_for_owner) is backup for the case that
+          # wrap_for_nesting fails, to ensure that the method will stil be defined in
+          # the correct place.
+          #
+          # @param [String] source  The source to wrap
+          # @return [String]
+          def wrap_for_owner(source)
+            Thread.current[:__pry_owner__] = @method.owner
+            source = "Thread.current[:__pry_owner__].class_eval do\n#{source}\nend"
+          end
+
+          # Update the new source code to have the correct Module.nesting.
+          #
+          # This method uses syntactic analysis of the original source file to determine
+          # the new nesting, so that we can tell the difference between:
+          #
+          #   class A; def self.b; end; end
+          #   class << A; def b; end; end
+          #
+          # The resulting code should be evaluated in the TOPLEVEL_BINDING.
+          #
+          # @param [String] source  The source to wrap.
+          # @return [String]
+          def wrap_for_nesting(source)
+            nesting = Pry::Code.from_file(@method.source_file).nesting_at(@method.source_line)
+
+            (nesting + [source] + nesting.map{ "end" } + [""]).join("\n")
+          rescue Pry::Indent::UnparseableNestingError => e
+            source
+          end
       end
 
       create_command(/amend-line(?: (-?\d+)(?:\.\.(-?\d+))?)?/) do

test/test_default_commands/test_introspection.rb

@@ -295,6 +295,15 @@ describe "Pry::DefaultCommands::Introspection" do
             def y?
               :because
             end
+
+            class B
+              G = :nawt
+
+              def foo
+                :maybe
+                G
+              end
+            end
           end
         EOS
         @tempfile.flush
@@ -403,6 +412,13 @@ describe "Pry::DefaultCommands::Introspection" do
           X.instance_method(:y?).owner.should == X
           X.new.y?.should == :maybe
         end
+
+        it "should preserve module nesting" do
+          mock_pry("edit-method -p X::B#foo")
+
+          X::B.instance_method(:foo).owner.should == X::B
+          X::B.new.foo.should == :nawt
+        end
       end
 
       describe 'on an aliased method' do