Take a limited number of items from a potentially infinite generator

Question

Hypothetically, I have a bunch of generators that never end, like:

function* progression() {
  let n = 0;

  while (true) {
    yield n++;
  }
}

I want to take first N items from that generator. So I made a function take for that:

function* take(n, generator) {
  let i = 0;
  let values = generator();

  for (let v of values) {
    if (++i > n) {
      return;
    }

    yield v;
  }
}

Is there a better way, maybe standard ES6+ API or a more concise syntax to achieve the same thing? If this function could fail under circumstances that I didn't account for, please give me a hint.

Answer 1

I would write it as:

function* take(n, generator) {
  let values = generator();

  let i = 0;
  while (i++ < n) {
    yield values.next().value;
  }
}

this seems to match the intent more closely though its mostly a matter of personal preference.

Update

You can also add the take method as an extension to the generator so you could call it like generator().take(3). See this SO answer if you are interested in that approach.

Answer 2

Given that the yielded value can be from N (1) to Infinity you can use .next(), check if the returned object property done is undefined, if true, break loop, else yield value, continue up to N.

A second function could also be used to pass a single or different N values and "a bunch of generators" to as rest parameters, where the return value will be an array of arrays or an object containing a property for each N and, or generator.

function* progression0() {
  let n = 0;
  while (true) {
    yield n++;
  }
}
function* progression1() {
  let n = 0;
  while (n < 1) {
    yield n++; // only 1 value yielded from this generator
  }
}
function* take(n, gen) {
  const g = gen();
  for (let i = 0; i < n; i++) {
    const {value, done} = g.next();
    // break if `done`
    if (done) {break};
    yield value;
  }
}
let n = 7;
let [...a] = take(n, progression0); // individual variable `a`
let [...b] = take(n, progression1); // individual variable `b`
console.log({a, b});
const processProgression = (n, ...gen) => gen.map(g => [...take(n, g)]);
let [c, d] = processProgression(n, progression0, progression1);
// same as `a`, `b` using function to process `N`, and n generators
console.log({c, d}); 

Answer 3

Some minor changes to improve use and protect misuse.

The take function only takes from the start. It would be more flexible if you pass the generator, not the generator function. That way you can take from a generator where you left off.

function *take(n, values) {    
    while (n-- > 0) { yield values.next().value }
}

and then used as

take(10, progression()); //to take the first 10

Or

const values = progression();
take(10, values); // take ten
take(10, values); // take ten more

There is a danger someone might use the infinite generator to fill an array

const array = [...progression()]; // uncatchable page crash

Which is an uncatchable error that will crash the page.

You could create an interface that protects the infinity by hiding the progression in closure. You can also then add some extras if need be.

const progression = (()=>{
    function *progression() {
        var n = 0;
        while (true) { yield n++ }
    }
    var current = progression();
    return {
        restart() { current = progression() },
        *take(count) {
            current = progression();
            while (count-- > 0) { yield current.next().value }
        },
        *more(count) {
            while (count-- > 0) { yield current.next().value }
        },
        get next() { return current.next().value }
    };
})();

And used as

progression.take(10); /// get first ten
progression.take(10); /// get first ten again
progression.more(10); /// get ten more
const v = progression.next;  // get next value (21st)
progression.restart();
const first = progression.next;  // get next value (1st)