Array#flatten in Ruby

gistfile1.txt

# Overwrites Array#flatten with a new pure Ruby implementation
class Array
  def flatten(depth=nil)
    ary = []
    self.each do |_|
      if _.is_a?(Array) && (depth.nil? || depth > 0)
        ary += _.flatten(depth.nil? ? nil : depth - 1)
      else
        ary << _
      end
    end
    ary
  end
  
  def flatten!
    replace flatten
  end
end

if __FILE__ == $0
  require 'test/unit'
  class TestArrayFlatten < Test::Unit::TestCase
    
    def test_1_dim_array
      assert_equal([1,2,"3",{4=>5},:'6'], 
                   [1,2,"3",{4=>5},:'6'].flatten)
    end

    def test_2_dim_array
      assert_equal([1,2,3,4,5,6], 
                   [1,2,[3,4,5],6].flatten)
    end
    
    def test_3_dim_array
      assert_equal([1,2,3,4,5,6], 
                   [1,2,[3,[4,5],6]].flatten)
    end

    def test_6_dim_array_with_depth_0
      assert_equal([1,[2,[3,[4,[5,[6]]]]]],
                   [1,[2,[3,[4,[5,[6]]]]]].flatten(0))
    end

    def test_6_dim_array_with_depth_1
      assert_equal([1,2,[3,[4,[5,[6]]]]],
                   [1,[2,[3,[4,[5,[6]]]]]].flatten(1))
    end

    def test_6_dim_array_with_depth_2
      assert_equal([1,2,3,[4,[5,[6]]]],
                   [1,[2,[3,[4,[5,[6]]]]]].flatten(2))
    end

    def test_6_dim_array_with_depth_3
      assert_equal([1,2,3,4,[5,[6]]],
                   [1,[2,[3,[4,[5,[6]]]]]].flatten(3))
    end
    
    def test_6_dim_array_with_depth_4
      assert_equal([1,2,3,4,5,[6]],
                   [1,[2,[3,[4,[5,[6]]]]]].flatten(4))
    end

    def test_6_dim_array_with_depth_5
      assert_equal([1,2,3,4,5,6],
                   [1,[2,[3,[4,[5,[6]]]]]].flatten(5))
    end

    def test_destructive
      ary = [1,2,[3,[4,5],6]]
      ary.flatten!
      assert_equal([1,2,3,4,5,6], ary)
    end
    
  end

end

So it won't flatten any other Enumerable then Array? How about adding a selection block:

  def flatten(depth=nil, &block)
    inject([]) { |ary, _| block.call(_) ? (ary += _.flatten) : (ary << _) }
  end

I did not implement depth, but that is part of 1.9 now. Also, inject isn't very efficient, so brute force it.

If you agree and polish this off, let me know as I would like to add to Facets.

Thanks for the suggestions; just removed the inject and it was about 40% faster. I personally can't see a big need for the block, but I also don't really convert Hashes into Arrays that often either.