Flatten arbitrarily nested collections.

flatten_enum.rb

class FlattenEnum
  class << self

    # Returns an enumerator that yields the non-enumerable items from within a
    # nested enumerable, in depth-first order.
    #
    # Example:
    #
    # FlattenEnum.create([[1, 2, [3]], 4]).to_a
    # #=> [1, 2, 3, 4]
    def create(nested_collection)
      Enumerator.new do |yielder|

        # The single-item case. Yield the item and stop the iteration.
        unless nested_collection.respond_to?(:each)
          item = nested_collection
          yielder << item
          raise StopIteration
        end

        # Recurse over each of the collection's elements in case they need
        # flattening.
        nested_collection.each do |element|
          inner_enum = create(element)
          loop do
            yielder << inner_enum.next
          end
        end

      end
    end

  end
end

flatten_test.rb

require 'minitest/autorun'

require_relative 'flatten_enum'

class TestFlatten < Minitest::Test
  def flatten(nested_collection)
    enum = FlattenEnum.create(nested_collection)
    enum.to_a
  end

  def test_single_integer
    input = [1]
    assert_equal([1], flatten(input))
  end

  def test_two_integers
    input = [2, 3]
    assert_equal([2, 3], flatten(input))
  end

  def test_single_nested_integer
    input = [[4]]
    assert_equal([4], flatten(input))
  end

  def test_full_customer_example
    input = [[1, 2, [3]], 4]
    expected = [1, 2, 3, 4]
    assert_equal(expected, flatten(input))
  end

  def nest_pathologically(item)
    [[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[ item ]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]
  end

  def test_pathological_depth
    input = nest_pathologically(
      [1, 2, nest_pathologically([3, nest_pathologically(4)])]
    )
    expected = [1, 2, 3, 4]
    assert_equal(expected, flatten(input))
  end

end