Flatten Exercise

flatten.rb

#!/usr/bin/ruby

# Ruby already has a 'flatten' method in the Array class
# [[1,2,[3]],4].flatten -> [1,2,3,4]
#
# I did my own implementation as an exercise

module Code
  def ruby_flatten(ary)
    ary.flatten if ary.respond_to? :flatten
  end

  def my_flatten(ary)
    _flatten = []

    ary.each do |element|
      _flatten += element.is_a?(Array) ? my_flatten(element) : [element]
    end

    _flatten
  end
end

class Test
  extend Code

  INPUT = [[1, 2, [3]], 4]
  EXPECTED = [1, 2, 3, 4]

  def self.ruby_flatten_test
    assertion = ruby_flatten(INPUT) == EXPECTED
    puts message("Ruby Flatten Test", assertion)
  end

  def self.my_flatten_test
    assertion = my_flatten(INPUT) == EXPECTED
    puts message("My Flatten Test", assertion)
  end

  private

  def self.message(test, assertion)
    "#{test}: #{assertion ? 'OK' : 'Failed'}"
  end
end

Test.ruby_flatten_test
Test.my_flatten_test