Tuesday, March 06, 2007

Enhance Array#collect to become magical

Do you do this often?

customers.collect { |customer| customer.name }
customers.collect { |customer| [customer.name, customer.id] }


Array#collect is indeed very powerful. But I still find myself to repeatedly declare a variable to keep a reference of the elements I am iterating. I could care less if I call it |customer| or |c|.

What if I enhance the Array class to do the following:

customers.collect_name
customers.collect_name_and_id


DRY-ness... Inspiration came from ActiveRecord's magic #find method.

class Array

  def method_missing(method_sym, *args)
    if collect_by_method?(method_sym)
      attributes = fetch_collect_attributes(method_sym)
      if attributes.size == 1
        block = lambda { |element| element.send attributes.first }
      else
        block = lambda { |element| attributes.collect { |attribute| element.send :"#{attribute}" } }
      end
      self.collect(&block)
    else
      super
    end
  end

  private

  def collect_by_method?(method_sym)
    method_sym.to_s =~ /^collect_/
  end

  def fetch_collect_attributes(method_sym)
    attributes = method_sym.to_s.gsub(/collect_/, '').split(/_and_/)
    raise ArgumentError, "Array#collect_* requires at least one method name after it. eg. #collect_id" if attributes.empty?
    attributes
  end

end


And of course, code is no good without tests:

class ArrayTest < Test::Unit::TestCase

  def test_collect_raises_exception_with_no_parameter
    assert_raise ArgumentError do
      [].collect_
    end
  end

  def test_collect_with_one_parameter
    array = []
    array << TestStruct.new(:id => 1, :foo => 'Foo 1')
    array << TestStruct.new(:id => 2, :foo => 'Foo 2')
    assert_equal ['Foo 1', 'Foo 2'], array.collect_foo
  end

  def test_collect_with_multiple_parameter
    array = []
    array << TestStruct.new(:id => 1, :foo => 'Foo 1', :bar => 'Bar 1')
    array << TestStruct.new(:id => 2, :foo => 'Foo 2', :bar => 'Bar 2')
    assert_equal [ ['Foo 1', 'Bar 1'], ['Foo 2', 'Bar 2'] ], array.collect_foo_and_bar
  end

  def test_collect_does_not_interfere_default_method_missing
    assert_raise NoMethodError do
    [].foo
    end
  end

end

2 comments:

Dr Nic said...

Hey, this looks similar to my map_by_method gem (articles, rubyforge). All good fun!

Dan Manges said...

I like this:
http://dev.rubyonrails.org/ticket/7878

collection.map(&[:attribute1, :attribute2])

No reliance on method missing...