Allow to query and replace subarrays of Rackstash::Fields::Array

lib/rackstash/fields/array.rb

@@ -52,27 +52,104 @@ module Rackstash
         RUBY
       end
 
-      # Retrieve a stored value from a given `index`
+      # Returns the element at `index`, or returns a subarray starting at the
+      # `start` index and continuing for `length` elements, or returns a subarray
+      # specified by `range` of indices.
+      #
+      # Negative indices count backward from the end of the array (-1 is the
+      # last element). For `start` and `range` cases the starting index is just
+      # before an element. Additionally, an empty array is returned when the
+      # starting index for an element range is at the end of the array.
+      #
+      # Returns `nil` if the index (or starting index) are out of range.
+      # @return [Object, nil]
+      #
+      # @overload [](index)
+      #   Returns the element at `index`
       #
       #   @param index [Integer] the index in the array where we fetch the value
-      # @return [Object, nil] the current value at `index` or `nil` if no value
-      #   could be found
-      def [](index)
-        @raw[index]
+      #   @return [Object, nil] the current value at `index` or `nil` if no
+      #     value could be found
+      #
+      # @overload [](start, length)
+      #   Returns an {Array} starting at the `start` index and continuing for
+      #   `length` elements
+      #
+      #   @param start [Integer] the index in the array where we fetch the
+      #     first value
+      #   @param length [Integer] the number of elements to return
+      #   @return [Array, nil] the current value at `index` or `nil` if `start`
+      #     is out of range
+      #
+      # @overload [](range)
+      #   Returns an {Array} starting at the `start` index and continuing for
+      #   `length` elements
+      #
+      #   @param range [Range] specifies the range of elements to return from
+      #     the array
+      #   @return [Array, nil] the current value at `index` or `nil` if the
+      #     start index is out of range
+      def [](index, length = nil)
+        result = length.nil? ? @raw[index] : @raw[index, length]
+        result = new(result) if ::Concurrent::Array === result
+        result
       end
+      alias slice []
 
-      # Set the value at a given index to the supplied value. The value is
-      # normalized before being set.
+      # Element Assignment - Sets the element at `index`, or replaces a subarray
+      # from the `start` index for `length` elements, or replaces a subarray
+      # specified by the `range` of indices.
       #
-      # You can set nested hashes and arrays here.
+      # All values are normalized before being set. You can set nested hashes
+      # and arrays here.
       #
+      # If indices are greater than the current capacity of the array, the array
+      # grows automatically. Elements are inserted into the array at start if
+      # length is zero.
+      #
+      # Negative indices will count backward from the end of the array. For
+      # `start` and `range` cases the starting index is just before an element.
+      #
+      # An `IndexError` is raised if a negative index points past the beginning
+      # of the array.
+      #
+      # See also {#push}, and {#unshift}.
+      #
+      # @overload []=(index, value)
       #   @param index [Integer] the index in the array where we set the value
       #   @param value [Object, Proc] any value which can be serialized to JSON.
       #     The value will be normalized before being set so that only JSON-
-      #   compatible objects are added into the array.
-      # @return [void]
-      def []=(index, value)
-        @raw[index] = normalize(value)
+      #     compatible objects are added into the array. A given Proc is called
+      #     with its result being used instead.
+      #
+      # @overload []=(range, value)
+      #   Replaces a subarray specified by the range of indices.
+      #
+      #   @param range [Range] the range if values in `self` which are replaced
+      #     by the passed `value`
+      #   @param value [Array, ::Array, #to_ary, Proc] An array contining
+      #     JSON-serializable values. The value will be normalized before being
+      #     set so that only JSON-compatible objects are added into the array. A
+      #     given Proc is called with its result being used instead.
+      #
+      # @overload []=(start, length, value)
+      #   Replaces a subarray from the `start` index for `length` elements with
+      #   the passed `value`.
+      #
+      #   @param index [Integer] the index in the array where we set the value
+      #   @param length [Integer] the index in the array where we set the value
+      #   @param value [Array, ::Array, #to_ary, Proc] An array contining
+      #     JSON-serializable values. The value will be normalized before being
+      #     set so that only JSON-compatible objects are added into the array. A
+      #     given Proc is called with its result being used instead.
+      #
+      # @return [value]
+      def []=(index, value_or_length, value = UNDEFINED)
+        if UNDEFINED.equal?(value)
+          @raw[index] = normalize(value_or_length)
+        else
+          @raw[index, value_or_length] = implicit(normalize(value, wrap: false))
+        end
       end
 
       # Add a given value at the end of the array

spec/rackstash/fields/array_spec.rb

@@ -64,14 +64,40 @@ describe Rackstash::Fields::Array do
   end
 
   describe '#[]' do
-    it 'returns nil if a value was not set' do
-      expect(array[1]).to be_nil
+    before do
+      array[0] = 'value'
+      array[1] = 'foo'
+      array[2] = 'bar'
+      array[3] = 'baz'
     end
 
     it 'returns a set value' do
-      array[0] = 'value'
       expect(array[0]).to eql 'value'
     end
+
+    it 'returns an array from start, end' do
+      expect(array[1, 3]).to be_a Rackstash::Fields::Array
+      expect(array[1, 3].to_ary).to eql %w[foo bar baz]
+
+      expect(array[2, 0].to_ary).to eql []
+      expect(array[2, 1].to_ary).to eql %w[bar]
+      expect(array[2, 5].to_ary).to eql %w[bar baz]
+    end
+
+    it 'returns an array from a range' do
+      expect(array[1..3]).to be_a Rackstash::Fields::Array
+      expect(array[1..3].to_ary).to eql %w[foo bar baz]
+
+      expect(array[2..4].to_ary).to eql %w[bar baz]
+      expect(array[2..-1].to_ary).to eql %w[bar baz]
+    end
+
+    it 'returns nil if a value was not set' do
+      expect(array[5]).to be_nil
+      expect(array[5, 2]).to be_nil
+      expect(array[2, -1]).to be_nil
+      expect(array[5..9]).to be_nil
+    end
   end
 
   describe '#[]=' do
@@ -81,6 +107,20 @@ describe Rackstash::Fields::Array do
       array[0] = 'value'
       expect(array[0]).to eql 'normalized'
     end
+
+    it 'can set values on a range' do
+      array.concat(%w[hello world with flowers and unicorns])
+
+      array[1..4] = %w[super duper]
+      expect(array.as_json).to eql %w[hello super duper unicorns]
+    end
+
+    it 'can set values from start, length' do
+      array.concat(%w[hello world with flowers and unicorns])
+
+      array[1, 4] = %w[shiny and sparkling]
+      expect(array.as_json).to eql %w[hello shiny and sparkling unicorns]
+    end
   end
 
   describe '#<<' do