tiny_hooks.rb in tiny_hooks-0.3.0

- old
+ new

@@ -6,18 +6,29 @@
 # `extend` this module and now you can define hooks with `define_hook` method.
 # See the test file for more detailed usage.
 module TinyHooks
   class Error < StandardError; end
 
+  # @api private
+  def self.extended(mod)
+    mod.class_eval { @@_originals ||= {} }
+    # mod.instance_variable_set(:@_originals, {}) unless mod.instance_variable_defined?(:@_originals)
+    # mod.define_singleton_method(:_originals) do
+    #   mod.instance_variable_get(:@_originals)
+    # end
+  end
+
   # Define hook with kind and target method
   #
   # @param [Symbol, String] kind the kind of the hook, possible values are: :before, :after and :around
   # @param [Symbol, String] target the name of the targeted method
   def define_hook(kind, target, &block)
     raise ArgumentError, 'You must provide a block' unless block
 
     original_method = instance_method(target)
+    @@_originals[target.to_sym] = original_method unless @@_originals[target.to_sym]
+
     body = case kind.to_sym
            when :before
              _before(original_method, &block)
            when :after
              _after(original_method, &block)
@@ -29,9 +40,19 @@
     undef_method(target)
     define_method(target, &body)
   end
 
   module_function :define_hook
+
+  # Restore original method
+  #
+  # @param [Symbol, String] target
+  def restore_original(target)
+    original_method = @@_originals[target.to_sym] || instance_method(target)
+
+    undef_method(target)
+    define_method(target, original_method)
+  end
 
   private
 
   def _before(original_method, &block)
     if RUBY_VERSION >= '2.7'