lib/remnant/discover.rb in remnant-0.2.4 vs lib/remnant/discover.rb in remnant-0.3.0

@@ -1,11 +1,14 @@
 class Remnant
   module Discover
     module ClassMethods
       def find(key, klass, method, instance = true)
         rediscover(key, klass, method, instance) if ActiveSupport::Dependencies.will_unload?(klass)
+        _inject(key, klass, method, instance)
+      end
 
+      def _inject(key, klass, method, instance)
         klass.class_eval <<-EOL, __FILE__, __LINE__
           #{"class << self" unless instance}
           def #{method}_with_remnant(*args, &block)
             ::Remnant::Discover.measure(#{key.inspect}) do
               #{method}_without_remnant(*args, &block)
@@ -15,10 +18,15 @@
           alias_method_chain :#{method}, :remnant
           #{"end" unless instance}
         EOL
       end
 
+      def find_with(klass, &block)
+        rediscover(klass, block) if ActiveSupport::Dependencies.will_unload?(klass)
+        block.call
+      end
+
       def measure(key, &block)
         if Remnant::Discover.running.include?(key)
           yield
         else
           result = nil
@@ -49,11 +57,16 @@
       def rediscover(*args)
         remnants_to_rediscover << args unless remnants_to_rediscover.include?(args)
       end
 
       def rediscover!
-        remnants_to_rediscover.map do |(key, klass_name, method, instance)|
-          find(key, klass_name.constantize, method, instance)
+        remnants_to_rediscover.map do |remnant|
+          if remnant.size == 4
+            key, klass_name, method, instance = remnant
+            _inject(key, klass_name.constantize, method, instance)
+          else
+            remant.first.call
+          end
         end
       end
     end
     extend ClassMethods
   end