lib/metacrunch/job/buffer.rb in metacrunch-4.0.3 vs lib/metacrunch/job/buffer.rb in metacrunch-4.1.0

@@ -1,26 +1,31 @@
 module Metacrunch
   class Job::Buffer
 
-    def initialize(size)
-      @size = size
+    def initialize(size_or_proc)
+      @size_or_proc = size_or_proc
+      @buffer = []
+
+      if @size_or_proc.is_a?(Numeric) && @size_or_proc <= 0
+        raise ArgumentError, "Buffer size must be a posive number greater that 0."
+      end
     end
 
     def buffer(data)
-      storage << data
-      flush if storage.count >= @size
+      @buffer << data
+
+      case @size_or_proc
+      when Numeric
+        flush if @buffer.count >= @size_or_proc.to_i
+      when Proc
+        flush if @size_or_proc.call == true
+      end
     end
 
     def flush
-      storage
+      @buffer
     ensure
-      @buffer = nil
-    end
-
-  private
-
-    def storage
-      @buffer ||= []
+      @buffer = []
     end
 
   end
 end