lib/json/streamer.rb in json-streamer-0.3.0 vs lib/json/streamer.rb in json-streamer-0.4.0

@@ -2,94 +2,112 @@
 require "json/stream"
 
 module Json
   module Streamer
     class JsonStreamer
+
+      attr_reader :aggregator
+
       def initialize(file_io, chunk_size = 1000)
         @parser = JSON::Stream::Parser.new
 
         @file_io = file_io
         @chunk_size = chunk_size
 
-        @object_nesting_level = 0
+        @current_nesting_level = -1
         @current_key = nil
         @aggregator = {}
         @temp_aggregator_keys = {}
 
         @parser.start_object {start_object}
         @parser.start_array {start_array}
         @parser.key {|k| key(k)}
-        @parser.value {|v| value(v)}
 
       end
 
-      def get_objects_from_level(yield_nesting_level)
-        @yield_nesting_level = yield_nesting_level
+      def get(nesting_level:-1, key:nil)
+        @yield_nesting_level = nesting_level
+        @wanted_key = key
 
+        @parser.value do |v|
+          if @aggregator[@current_nesting_level].kind_of? Array
+            @aggregator[@current_nesting_level] << v
+          else
+            @aggregator[@current_nesting_level][@current_key] = v
+            if yield_value?
+              yield v
+            end
+          end
+        end
+
         # Callback containing yield has be defined in the method called via block
         @parser.end_object do
-          if @object_nesting_level.eql? @yield_nesting_level
-            yield @aggregator[@object_nesting_level].clone
+          if yield_object?
+            yield @aggregator[@current_nesting_level].clone
             # TODO probably can be faster than reject!{true}
-            @aggregator[@object_nesting_level].reject!{true}
+            @aggregator[@current_nesting_level].reject!{true}
           else
             merge_up
           end
 
-          @object_nesting_level -= 1
+          @current_nesting_level -= 1
         end
 
         @parser.end_array do
-          if @object_nesting_level.eql? @yield_nesting_level
-            yield @aggregator[@object_nesting_level].clone
+          if yield_object?
+            yield @aggregator[@current_nesting_level].clone
             # TODO probably can be faster than reject!{true}
-            @aggregator[@object_nesting_level].reject!{true}
+            @aggregator[@current_nesting_level].reject!{true}
           else
             merge_up
           end
 
-          @object_nesting_level -= 1
+          @current_nesting_level -= 1
         end
 
         @file_io.each(@chunk_size) do |chunk|
           @parser << chunk
         end
       end
 
+      def yield_object?
+        @current_nesting_level.eql? @yield_nesting_level or (not @wanted_key.nil? and @wanted_key == @temp_aggregator_keys[@current_nesting_level-1])
+      end
+
+      def yield_value?
+        @wanted_key == @current_key
+      end
+
       def start_object
-        @temp_aggregator_keys[@object_nesting_level] = @current_key
-        @object_nesting_level += 1
-        @aggregator[@object_nesting_level] = {}
+        @temp_aggregator_keys[@current_nesting_level] = @current_key
+        @current_nesting_level += 1
+        @aggregator[@current_nesting_level] = {}
       end
 
       def start_array
-        @temp_aggregator_keys[@object_nesting_level] = @current_key
-        @object_nesting_level += 1
-        @aggregator[@object_nesting_level] = []
+        @temp_aggregator_keys[@current_nesting_level] = @current_key
+        @current_nesting_level += 1
+        @aggregator[@current_nesting_level] = []
       end
 
       def key k
         @current_key = k
       end
 
       def value v
-        if @aggregator[@object_nesting_level].kind_of? Array
-          @aggregator[@object_nesting_level] << v
-        else
-          @aggregator[@object_nesting_level][@current_key] = v
-        end
+
       end
 
       def merge_up
-        return if @object_nesting_level == 1
-        previous_object_nesting_level = @object_nesting_level - 1
-        if @aggregator[previous_object_nesting_level].kind_of? Array
-          @aggregator[previous_object_nesting_level] << @aggregator[@object_nesting_level]
+        return if @current_nesting_level == 0
+        previous_nesting_level = @current_nesting_level - 1
+        if @aggregator[previous_nesting_level].kind_of? Array
+          @aggregator[previous_nesting_level] << @aggregator[@current_nesting_level]
         else
-          @aggregator[previous_object_nesting_level][@temp_aggregator_keys[previous_object_nesting_level]] = @aggregator[@object_nesting_level]
+          @aggregator[previous_nesting_level][@temp_aggregator_keys[previous_nesting_level]] = @aggregator[@current_nesting_level]
         end
 
-        @aggregator.delete(@object_nesting_level)
+        @aggregator.delete(@current_nesting_level)
         @aggregator
       end
     end
   end
 end