Implementing Lazy take
Now that you’ve implemented map, you can have a go at implementing the take method. As its name suggests, Enumerable#take(n) returns the first n elements from the Enumerable. As with the lazy version of map, the lazy version of take also returns a Lax instance, this time wrapping the Enumerable#take method. Here’s how it looks:
| | def take(n) |
| | taken = 0 |
| | Lax.new(self) do |yielder, val| |
| | if taken < n |
| | yielder << val |
| | taken += 1 |
| | else |
| | raise StopIteration |
| | end |
| | end |
| | end |
The logic for take should be easy enough for you to follow. The interesting thing here is how take signals that the iteration has ended. When taken reaches the limit, a StopIteration exception is raised to break out of the block.
While the use of exceptions for control flow are generally frowned upon, this is exactly how Ruby implements take. take throws an exception once the enumerator takes the right number of elements. That exception is handled inside the constructor.
What should be done once you get all the values you need? Well, nothing much really. All you need is to handle the exception silently in the initialize method:
| | def initialize(receiver) |
| | super() do |yielder| |
| » | begin |
| | receiver.each do |val| |
| | if block_given? |
| | yield(yielder, val) |
| | else |
| | yielder << val |
| | end |
| | end |
| » | rescue StopIteration |
| » | end |
| | end |
| | end |
Before you take the code for another spin, here’s what you should have so far:
| | module Enumerable |
| | def lax |
| | Lax.new(self) |
| | end |
| | end |
| | |
| | class Lax < Enumerator |
| | |
| | def initialize(receiver) |
| | super() do |yielder| |
| | begin |
| | receiver.each do |val| |
| | if block_given? |
| | yield(yielder, val) |
| | else |
| | yielder << val |
| | end |
| | end |
| | rescue StopIteration |
| | end |
| | end |
| | end |
| | |
| | def map(&block) |
| | Lax.new(self) do |yielder, val| |
| | yielder << block.call(val) |
| | end |
| | end |
| | |
| | def take(n) |
| | taken = 0 |
| | Lax.new(self) do |yielder, val| |
| | if taken < n |
| | yielder << val |
| | taken += 1 |
| | else |
| | raise StopIteration |
| | end |
| | end |
| | end |
| | |
| | end |
With that in place, let’s try out the code:
| | >> 1.upto(Float::INFINITY).lax.map { |x| x*x }.map { |x| x+1 }.take(5).to_a |
| | => [2, 5, 10, 17, 26] |
If you get [2, 5, 10, 17, 26], then you should pat yourself on the back. If not, go back and review your implementation carefully.