When to Define Your Own Custom Ruby Iterator?

General rule of thumb is if you have a collection of objects, you should build custom iterators so that users could readily use them like they do with other collection such as Array, Hash, File, etc.

The easiest way to do this is to mix-in the Enumerable module. The only requirement is to define an each method in your class which will provide Enumerable with your collection objects.

Let’s take a look at an example. Let’s say you’re a property owner with a few properties in various cities.

class Property
  attr_reader :address, :sq_ft, :num_of_rooms

  def initialize(address, sq_ft, num_of_rooms)
    @address = address
    @sq_ft = sq_ft
    @num_of_rooms = num_of_rooms
  end

  def get_info
    puts "Property is located at #{address} and is #{sq_fit} sq/ft with #{num_of_rooms} rooms"
  end
end

prop1 = Property.new("123 Westington St.", 1000, 2)
prop2 = Property.new("322 Popular St.", 2000, 3)
prop3 = Property.new("932 Horray St.", 1000, 7)
prop4 = Property.new("213 Hello Hat St.", 9000, 9)
prop5 = Property.new("858 Bacon St.", 3000, 3)
prop6 = Property.new("391 Blahestington St.", 10000, 3)

class PropertyList
  attr_reader :name, :property_list

  def initialize(name, property_list=[])
    @name = name
    @property_list = property_list
  end

  def add_property(p)
    self.property_list << p
  end
end

la_props = PropertyList.new("Los Angeles")
la_props.add_property(prop2)
la_props.add_property(prop4)
la_props.add_property(prop6)

Now, what if we want to iterate through the addresses in LA?

la_props.each { |p| puts p.address }

# undefined method `each' for #<PropertyList:0x007fdbac0ab998> (NoMethodError)

So how do we implement an iterator, much like we’ve seen with other collections such as Arrays and Hashes? We define an each method in our PropertyList class.

  def each
    self.property_list.each { |p| yield p }
  end

# Output:
# 322 Popular St.
# 213 Hello Hat St.
# 391 Blahestington St.

Now what if we want to get the total square footage of all the properties in LA that have at more than 3 rooms? Naturally we think of the select and reduce methods, but those belong to the Enumerable module.

la_sq_ft = la_props.select { |s| s.num_of_rooms > 3 }.map { |s| s.sq_ft }.reduce(0, :+)
puts "Total SQ/FT for LA Properties: #{la_sq_ft}"

# private method `select' called for #<PropertyList:0x007feb0e03a010> (NoMethodError)

Aha! As mentioned, all we have to do is define an each method in the host class to mix in Enumerable. Our final class looks like this:

class PropertyList
  include Enumerable

  attr_reader :name, :property_list

  def initialize(name, property_list=[])
    @name = name
    @property_list = property_list
  end

  def add_property(p)
    self.property_list << p
  end

  def each
    self.property_list.each { |p| yield p }
  end
end

and we get an output of: “Total SQ/FT for LA Properties: 10000” which is correct.

Now let’s level up some more. How about we build our own Enumerable module. We want to get this to work:

big_rooms      = la_props.select { |s| s.num_of_rooms > 3 }
only_sq_ft     = la_props.map { |s| s.sq_ft }
sum_all_sq_ft  = la_props.reduce(:+)

All we have to do is define our own module and include it in the PropertyList class:

module EstateEnumerable
  # Here we would expect an each method in the class
  def select
    puts "Using custom select!"
    result = []
    each do |s|
      result << s if yield s
    end
    result
  end

  def map
    puts "Using custom map!"
    result = []
    each do |m|
      r = yield m
      result << r
    end
    result
  end

  def reduce(op)
    puts "Using custom reduce!"
    case op
    when Symbol
      curr_proc = op.to_proc
      result = 0
      each do |o|
        result = curr_proc.call(result, o.sq_ft)
      end
      result
    else
        puts "Need a symbol"
    end
  end
end

With this we just “Include EstateEnumerable” instead of “Enumerable” and we’ll get this as the output:

puts "Plenty of rooms: #{big_rooms}"
puts "Filter only for sq ft: #{only_sq_ft}"
puts "Sum of all sq ft: #{sum_all_sq_ft}"

# Plenty of rooms: [#<Property:0x007fa0a9039e08 @address="932 Horray St.", @sq_ft=1000, @num_of_rooms=7>, #<Property:0x007fa0a9039d90 @address="213 Hello Hat St.", @sq_ft=9000, @num_of_rooms=9>]
# Filter only for sq ft: [1000, 2000, 1000, 9000, 3000, 10000]
# Sum of all sq ft: 26000

This hopefully gives you a glimpse of what Ruby blocks are capable of.