diff --git a/packages/algorithm/src/take.ts b/packages/algorithm/src/take.ts index f62b3a302..bb12709aa 100644 --- a/packages/algorithm/src/take.ts +++ b/packages/algorithm/src/take.ts @@ -35,9 +35,12 @@ export function* take( object: Iterable, count: number ): IterableIterator { - for (const value of object) { - if (0 < count--) { - yield value; - } + if (count < 1) { + return; + } + const it = object[Symbol.iterator](); + let item: IteratorResult; + while (0 < count-- && !(item = it.next()).done) { + yield item.value; } } diff --git a/packages/algorithm/tests/src/take.spec.ts b/packages/algorithm/tests/src/take.spec.ts index 83a3828c7..fe96d2ca4 100644 --- a/packages/algorithm/tests/src/take.spec.ts +++ b/packages/algorithm/tests/src/take.spec.ts @@ -20,13 +20,17 @@ describe('@lumino/algorithm', () => { describe('take() from an iterable iterator', () => { testIterator(() => { - return [take([0, 1, 2, 3][Symbol.iterator](), 1), [0]]; + return [take([0, 1, 2, 3], 1), [0]]; }); }); describe('take() with count=0', () => { - testIterator(() => { - return [take([0, 1, 2, 3][Symbol.iterator](), 0), []]; - }); + testIterator(() => [take([0, 1, 2, 3], 0), []]); + }); + + describe('take() only takes as many as count', () => { + const it = [0, 1, 2, 3][Symbol.iterator](); + testIterator(() => [take(it, 2), [0, 1]]); + testIterator(() => [take(it, 4), [2, 3]]); }); });