Lazily adding information to Ruby's Enumerator's.

enumerator_with_final.rb

module Enumerator::WithFinal
  # Emits a boolean indicating whether or not the current iteration is the final one.
  def self.with_final(prev_enum)
    begin
      while next_val = prev_enum.next
        is_final = begin
           prev_enum.peek
        rescue StopIteration
          true
        else
          false
        end
        yield [next_val, is_final]
      end
    rescue StopIteration
    end
  end

  # If you like monkeypatching Ruby's standard library, go for it.
  def with_final
    Enumerator.new {|y| Enumerator::WithFinal.with_final(self) {|*v| y << v } }
  end
end

if __FILE__ == $0
  require "minitest/spec"
  require "minitest/autorun"

  describe Enumerator::WithFinal do
    it "chains well with others" do
      r = 1..10
      Enumerator::WithFinal.with_final(r.each.with_object("foo").with_index) do |*args|
        val, foo, i, is_final = args.flatten
        # puts "val, i, foo, is_final: #{val}, #{i}, #{foo}, #{is_final}"
        assert_equal(val, i + 1)
        assert_equal(is_final, i == r.end - 1)
      end
    end

    it "monkeypatches" do
      class Enumerator
        include Enumerator::WithFinal
      end
      r = 1..10
      r.each.with_object("foo").with_index.with_final.each do |*args|
        val, foo, i, is_final = args.flatten
        # puts "val, i, foo, is_final: #{val}, #{i}, #{foo}, #{is_final}"
        assert_equal(val, i + 1)
        assert_equal(is_final, i == r.end - 1)
      end
    end
  end
end