How to count how many times an iterator will call yield?

Question

I have a method, foo, that yields objects. I want to count the number of objects it yields.

I have

def total_foo
  count = 0
  foo { |f| count += 1}
  count
end

but there's probably a better way. Any ideas for this new Rubyist?

Here's the definition for foo (it's a helper method in Rails):

def foo(resource=@resource)
  resource.thingies.each do |thingy|
    bar(thingy) { |b| yield b }  # bar also yields objects
  end
end

Answer 1

Any method that calls yield can be used to build an Enumerator object, on which you can call count, by means of the Object#to_enum method. Remember that when you call count the iterator is actually executed so it should be free of side effects! Following a runnable example that mimics your scenario:

@resources = [[1,2], [3,4]]

def foo(resources = @resources)
  resources.each do |thingy|
    thingy.each { |b| yield b }
  end
end

foo { |i| puts i }
# Output:
# 1
# 2
# 3
# 4

to_enum(:foo).count
# => 4

You can pass an argument to foo:

to_enum(:foo, [[5,6]]).count
# => 2

Alternatively you can define foo to return an Enumerator when it's called without a block, this is the way stdlib's iterators work:

def foo(resources = @resources)

  return to_enum(__method__, resources) unless block_given?

  resources.each do |thingy|
    thingy.each { |b| yield b }
  end
end

foo.count
# => 4

foo([[1,2]]).count
# => 2

foo([[1,2]]) { |i| puts i }
# Output:
# 1
# 2

You can pass a block to to_enum that is called when you call size on the Enumerator to return a value:

def foo(resources = @resources)
  unless block_given?
    return to_enum(__method__, resources) do
      resources.map(&:size).reduce(:+) # thanks to @Ajedi32
    end
  end

  resources.each do |thingy|
    thingy.each { |b| yield b }
  end
end

foo.size
# => 4

foo([]).size
# => 0

In this case using size is sligthly faster than count, your mileage may vary.

Answer 2

Assuming you otherwise only care about the side-effect of foo, you could have foo itself count the iterations:

def foo(resource=@resource)
  count = 0
  resource.thingies.each do |thingy|
    bar(thingy) do |b|
      count += 1
      yield b
    end  # bar also yields objects
  end
  count
end

And then:

count = foo { |f| whatever... }

You can also ignore the return value if you choose, so just:

foo { |f| whatever... }

In cases you don't care what the count is.

There may be better ways to handle all of this depending upon the bigger context.