added before_command and after_command decorators + tests

lib/pry/command_set.rb

@@ -142,6 +142,43 @@ class Pry
       commands[name] = Command.new(name, description, options, block)
     end
 
+
+    # Execute a block of code before a command is invoked. The block also
+    # gets access to parameters that will be passed to the command and
+    # is evaluated in the same context.
+    # @param [String, Regexp] name The name of the command.
+    # @yield The block to be run before the command.
+    # @example Display parameter before invoking command
+    #   Pry.commands.before_command("whereami") do |n|
+    #     output.puts "parameter passed was #{n}"
+    #   end
+    def before_command(name, &block)
+      prev_block = commands[name].block
+      wrapper_block = proc do |*args|
+        instance_exec(*args, &block)
+        instance_exec(*args, &prev_block)
+      end
+      commands[name].block = wrapper_block
+    end
+
+    # Execute a block of code after a command is invoked. The block also
+    # gets access to parameters that will be passed to the command and
+    # is evaluated in the same context.
+    # @param [String, Regexp] name The name of the command.
+    # @yield The block to be run after the command.
+    # @example Display text 'command complete' after invoking command
+    #   Pry.commands.after_command("whereami") do |n|
+    #     output.puts "command complete!"
+    #   end
+    def after_command(name, &block)
+      prev_block = commands[name].block
+      wrapper_block = proc do |*args|
+        instance_exec(*args, &prev_block)
+        instance_exec(*args, &block)
+      end
+      commands[name].block = wrapper_block
+    end
+
     def each &block
       @commands.each(&block)
     end

test/test_command_set.rb

@@ -108,7 +108,7 @@ describe Pry::CommandSet do
     run.should == true
   end
 
-  it 'should be able to change the descritpions of methods' do
+  it 'should be able to change the descriptions of methods' do
     @set.command('foo', 'bar') {}
     @set.desc 'foo', 'baz'
 
@@ -226,4 +226,95 @@ describe Pry::CommandSet do
 
     order.should == order.sort
   end
+
+  describe "command decorators - before_command and after_command" do
+    describe "before_command" do
+      it 'should be called before the original command' do
+        foo = []
+        @set.command('foo') { foo << 1 }
+        @set.before_command('foo') { foo << 2 }
+        @set.run_command(@ctx, 'foo')
+
+        foo.should == [2, 1]
+      end
+
+      it 'should share the context with the original command' do
+        @ctx.target = "test target string"
+        before_val  = nil
+        orig_val    = nil
+        @set.command('foo') { orig_val = target }
+        @set.before_command('foo') { before_val = target }
+        @set.run_command(@ctx, 'foo')
+
+        before_val.should == @ctx.target
+        orig_val.should == @ctx.target
+      end
+
+      it 'should work when applied multiple times' do
+        foo = []
+        @set.command('foo') { foo << 1 }
+        @set.before_command('foo') { foo << 2 }
+        @set.before_command('foo') { foo << 3 }
+        @set.before_command('foo') { foo << 4 }
+        @set.run_command(@ctx, 'foo')
+
+        foo.should == [4, 3, 2, 1]
+      end
+
+    end
+
+    describe "after_command" do
+      it 'should be called after the original command' do
+        foo = []
+        @set.command('foo') { foo << 1 }
+        @set.after_command('foo') { foo << 2 }
+        @set.run_command(@ctx, 'foo')
+
+        foo.should == [1, 2]
+      end
+
+      it 'should share the context with the original command' do
+        @ctx.target = "test target string"
+        after_val  = nil
+        orig_val    = nil
+        @set.command('foo') { orig_val = target }
+        @set.after_command('foo') { after_val = target }
+        @set.run_command(@ctx, 'foo')
+
+        after_val.should == @ctx.target
+        orig_val.should == @ctx.target
+      end
+
+      it 'should determine the return value for the command' do
+        @set.command('foo', 'bar', :keep_retval => true) { 1 }
+        @set.after_command('foo') { 2 }
+        @set.run_command(@ctx, 'foo').should == 2
+      end
+
+      it 'should work when applied multiple times' do
+        foo = []
+        @set.command('foo') { foo << 1 }
+        @set.after_command('foo') { foo << 2 }
+        @set.after_command('foo') { foo << 3 }
+        @set.after_command('foo') { foo << 4 }
+        @set.run_command(@ctx, 'foo')
+
+        foo.should == [1, 2, 3, 4]
+      end
+    end
+
+    describe "before_command and after_command" do
+      it 'should work when combining both before_command and after_command' do
+        foo = []
+        @set.command('foo') { foo << 1 }
+        @set.after_command('foo') { foo << 2 }
+        @set.before_command('foo') { foo << 3 }
+        @set.run_command(@ctx, 'foo')
+
+        foo.should == [3, 1, 2]
+      end
+
+    end
+
+  end
 end