How can I create an enumerator that does certain things after iteration?

Question

How can I create an enumerator that optionally takes a block? I want some method foo to be called with some arguments and an optional block. If it is called with a block, iteration should happen on the block, and something should be done on it, involving the arguments given. For example, if foo were to take a single array argument, apply map on it with the given block, and return the result of join applied to it, then it would look like:

foo([1, 2, 3]){|e| e * 3}
# => "369"

If it is called without a block, it should return an enumerator, on which instance methods of Enumerator (such as with_index) should be able to apply, and execute the block in the corresponding way:

enum = foo([1, 2, 3])
# => Enumerator
enum.with_index{|e, i| e * i}
# => "026"

I defined foo using a condition to see if a block is given. It is easy to implement the case where the block is given, but the part returning the enumerator is more difficult. I guess I need to implement a sublass MyEnum of Enumerator and make foo return an instance of it:

def foo a, &pr
  if pr
    a.map(&pr).join
  else
    MyEnum.new(a)
  end
end

class MyEnum < Enumerator
  def initialize a
    @a = a
    ...
  end
  ...
end

But calling MyEnum.new raises a warning message: Enumerator.new without a block is deprecated; use Object#to_enum. If I use to_enum, I think it would return a plain Enumerator instance, not the MyEnum with the specific feature built in. On top of that, I am not sure how to implement MyEnum in the first place. How can I implement such enumerator? Or, what is the right way to do this?

Answer 1

You could do something like this.

def foo a, &pr
  if pr
    a.map(&pr).join
  else
    o = Object.new
    o.instance_variable_set :@a, a
    def o.each *y
      foo @a.map { |z| yield z, *y } { |e| e }
    end
    o.to_enum
  end
end

Then we have

enum = foo([1,2,3])
enum.each { |x| 2 * x } # "246"

or

enum = foo([1,2,3])
enum.with_index { |x, i| x * i } # "026"

Inspiration was drawn from the Enumerator documentation. Note that all of your expectations about enumerators like you asked for hold, because .to_enum takes care of all that. enum is now a legitimate Enumerator!

enum.class # Enumerator