My own implementation of Array#flatten

array_flatten.rb

class ArrayFlatten
  def initialize(arr)
    @arr = arr
  end

  def flatten(level = nil)
    _flatten(@arr, 0, level)
  end

  private

  # _flatten method receives an array and returns a new array flattened (one dimension)
  # this method could receive an optional argument called level, indicating
  # level of recursion for the new array. level should be a positive integer

  def _flatten(arr, current, level = nil, flattened_arr = [])
    arr.each do |elem|
      if elem.is_a?(Array) && (level.nil? || current < level)
        _flatten(elem, current + 1, level, flattened_arr)
      else
        flattened_arr << elem
      end
    end
    flattened_arr
  end
end

array_flatten_spec.rb

require 'rspec'
require_relative 'array_flatten'

RSpec.describe ArrayFlatten do
  describe '#flatten' do
    context 'when method no takes argument' do
      let(:arr_one) { ArrayFlatten.new([1, 2, [3, 4]]) }
      let(:arr_two) { ArrayFlatten.new([1, 2, [3, [4, 5, 6]]]) }

      it 'returns flattened array' do
        expect(arr_one.flatten).to eq [1, 2, 3, 4]
      end

      it 'returns a array different than original array' do
        expect(arr_two.flatten).not_to eq arr_two
      end
    end

    context 'when method receives level argument' do
      let(:arr_three) { ArrayFlatten.new([1, [2, [3, [4, [5]]]]]) }
      let(:arr_four) { ArrayFlatten.new([[[[1, 2], [3, [4, 5, 6]]]]]) }

      it 'returns flattened array until level 1' do
        expect(arr_three.flatten(1)).to eq [1, 2, [3, [4, [5]]]]
      end

      it 'returns flattened array until level 3' do
        expect(arr_four.flatten(3)).to eq [1, 2, 3, [4, 5, 6]]
      end
    end
  end
end