Collection classes in Ruby


Iterating over collections in Ruby is fun. Methods like each, map, and inject are intuitive and easy to use. When I find myself doing the same iteration logic over and over, though, I get frustrated. Duplicate iteration logic is a sign of poor design, and loop-happy methods can be harder to read and harder to test.

To avoid these issues, I’ve started creating Collection classes to wrap collections of objects. For example, say we have a Document class that looks like this:

class Document
  attr_reader :filename

  def initialize filename
    @filename = filename
  end

  def read
    @content ||= File.open(filename).read
  end
end

Suppose a user can have many documents. Let’s write a method to read all documents for a user:

class User
  def documents
    # returns an array of all the user's Document objects
  end

  def read_documents 
    self.documents.map do |doc|
      doc.read
    end.join
  end
end

While this does the job, what if we have a Team class that can also have documents? If we want a read_documents method for Team, we’ll be rewriting the exact same loop.

Let’s see how a DocumentCollection class could help us avoid this duplication.

Creating the Collection

First, we create a simple class to wrap a collection of documents:

class DocumentCollection
  def initialize documents=[]
    @documents = documents
  end
end

With our base class at hand, let’s add a read_all method to mimic our User#read_documents method.

class DocumentCollection
  def read_all
    @documents.map do |doc|
      doc.read
    end.join
  end
end

Great! Now we have a read_all method defined in a single place. Now we can update our User#read_documents object to use the new collection, like so:

class User
  def read_documents
    DocumentCollection.new(self.documents).read_all
  end
end

Adding this logic to a Team class is now easy, and we are only defining the read_all logic in one place. So we’re done, right?

Not exactly. Although our DocumentCollection class meets our current needs, how easy will it be to extend? For example, finding documents of a certain type might mean defining a matching_type method:

  def matching_type type
    @documents.select do |doc|
      doc.type == type
    end
  end

This works fine, but notice how all our methods are just loops on our @documents collection. How could we clean this up? By using Ruby’s wonderful Enumerable module, that’s how!

Improving with Enumerable

Including Enumerable in a class gives you all the iterative powers of classes like Array without any extra work. All we have to do is define an each method to tell Enumerable how we want to iterate over our class.

Here’s how our DocumentCollection class might look after including Enumerable:

class DocumentCollection
  include Enumerable

  def initialize documents=[]
    @documents = documents
  end

  def each &block
    @documents.each(&block)
  end

  def read_all
    map {|doc| doc.read}.join
  end

  def matching_type type
    select {|doc| doc.type == type}
  end
end

See how much cleaner that is than our initial version? Thanks to Enumerable, adding iteration logic is a snap. Now we can encapsulate any collection logic (and testing) in our collection class and share it across our codebase. Oh, the magic of Enumerable!