DIY Array#flatten

array_extensions.rb

module ArrayExtensions
  refine Array do
    def my_flatten
      flat_enumerator.to_a
    end

    private

    def flat_enumerator
      Enumerator.new do |out|
        visitor = -> (item) { item.respond_to?(:to_ary) ? item.to_ary.each(&visitor) : out << item }
        visitor.call(self)
      end
    end
  end
end

array_extensions_spec.rb

require_relative "./array_extensions"

using ArrayExtensions

describe ArrayExtensions do
  describe "#my_flatten" do
    it "flattens empty array" do
      expect([].my_flatten).to eq([])
    end

    it "flattens 1-level array" do
      expect([1, 2, 3].my_flatten).to eq([1, 2, 3])
    end

    it "flattens nested array" do
      expect([1, [2, 3]].my_flatten).to eq([1, 2, 3])
    end

    it "flattens deep-nested array" do
      expect([1, [2, 3, [4]]].my_flatten).to eq([1, 2, 3, 4])
    end

    it "is does not try to iterate through Hashes or Strings" do
      expect(["hello", {2 => 3}, 4].my_flatten).to eq(["hello", {2 => 3}, 4])
    end
  end
end