lib/rib/shell.rb in rib-1.5.2 vs lib/rib/shell.rb in rib-1.5.3

@@ -2,11 +2,10 @@
 require 'rib/plugin'
 require 'rib/api'
 
 module Rib; class Shell
   include API
-  trap('INT'){ raise Interrupt }
 
   def self.use mod
     include mod
   end
 
@@ -45,19 +44,20 @@
   end
 
   # Loops shell until user exits
   def loop
     before_loop
+    set_trap
     start
     in_loop
     stop
     self
   rescue Exception => e
     Rib.warn("Error while running loop:\n  #{format_error(e)}")
     raise
   ensure
-
+    release_trap
     after_loop
   end
 
   def start
     @running = true
@@ -75,9 +75,17 @@
     @warnings ||= []
   end
 
   protected
   attr_writer :config
+
+  def set_trap
+    @trap_proc = trap('INT'){ raise Interrupt }
+  end
+
+  def release_trap
+    trap('INT', &@trap_proc) if @trap_proc.kind_of?(Proc)
+  end
 
   private
   # Avoid namespace pollution from rubygems bin stub.
   # To be specific, version and str.
   def new_private_binding