Add product(...) and repeat(x, n) support

Author: blakeembrey

README.md

@@ -36,7 +36,7 @@ Make an iterator returning elements from the iterable and saving a copy of each.
 cycle([1, 2, 3]); //=> [1, 2, 3, 1, 2, 3, 1, 2, ...]
 ```
 
-### `repeat<T>(value: T): Iterable<T>`
+### `repeat<T>(value: T, times?: number): Iterable<T>`
 
 Make an iterator that repeats `value` over and over again.
 
@@ -271,6 +271,14 @@ Sums `start` and the items of an `iterable` from left to right and returns the t
 sum([1, 2, 3, 4, 5]); //=> 15
 ```
 
+### `product<T>(...iterables: Iterable<T>[]): Iterable<T[]>`
+
+Cartesian product of input iterables.
+
+```ts
+product("ABCD", "xy"); //=> Ax Ay Bx By Cx Cy Dx Dy
+```
+
 ## Reference
 
 - [Itertools Recipes](https://docs.python.org/3/library/itertools.html#itertools-recipes)

src/async.ts

@@ -72,13 +72,11 @@ export function contains<T>(iterable: AnyIterable<T>, needle: T) {
 /**
  * Returns an iterable of enumeration pairs.
  */
-export async function* enumerate<T>(
+export function enumerate<T>(
   iterable: AnyIterable<T>,
   offset = 0
 ): AsyncIterable<[number, T]> {
-  let index = offset;
-
-  for await (const value of iterable) yield [index++, value];
+  return zip(range(offset), iterable);
 }
 
 /**
@@ -188,8 +186,9 @@ export async function* cycle<T>(iterable: AnyIterable<T>): AsyncIterable<T> {
 /**
  * Make an iterator that repeats `value` over and over again.
  */
-export async function* repeat<T>(value: T): AsyncIterable<T> {
-  while (true) yield value;
+export async function* repeat<T>(value: T, times?: number): AsyncIterable<T> {
+  if (times === undefined) while (true) yield value;
+  for (let i = 0; i < times; i++) yield value;
 }
 
 /**
@@ -642,3 +641,37 @@ export async function sum(
 ): Promise<number> {
   return reduce(iterable, (x, y) => x + y, start);
 }
+
+/**
+ * Recursively produce all produces of a list of iterators.
+ */
+async function* _product<T>(
+  pools: AnyIterator<T | typeof SENTINEL>[],
+  buffer: T[] = []
+): AsyncIterable<T[]> {
+  if (pools.length === 0) {
+    yield buffer.slice();
+    return;
+  }
+
+  const [pool, ...others] = pools;
+
+  while (true) {
+    const item = await pool.next();
+    if (item.value === SENTINEL) break;
+    buffer.push(item.value);
+    yield* _product(others, buffer);
+    buffer.pop();
+  }
+}
+
+/**
+ * Cartesian product of input iterables.
+ */
+export async function* product<T extends any[]>(
+  ...iterables: AnyTupleIterable<T>
+): AsyncIterable<T> {
+  const pools = iterables.map(x => iter(cycle(chain(x, repeat(SENTINEL, 1)))));
+
+  yield* _product(pools) as AsyncIterable<T>;
+}

src/index.spec.ts

@@ -149,6 +149,12 @@ describe("iterative", () => {
 
       expect(Array.from(iterable)).toEqual([10, 10, 10, 10, 10]);
     });
+
+    it("should repeat a value up to times", () => {
+      const iterable = iter.repeat(10, 5);
+
+      expect(Array.from(iterable)).toEqual([10, 10, 10, 10, 10]);
+    });
   });
 
   describe("cycle", () => {
@@ -469,8 +475,28 @@ describe("iterative", () => {
     it("should sum an iterable", () => {
       expect(iter.sum(iter.range(0, 10))).toEqual(45);
     });
+
     it("should sum an iterable with custom start", () => {
       expect(iter.sum(iter.range(0, 10), 5)).toEqual(50);
     });
   });
+
+  describe("product", () => {
+    it("should generate the product of multiple iterators", () => {
+      const iterable = iter.product("ABCD", "xy");
+
+      const result = [
+        ["A", "x"],
+        ["A", "y"],
+        ["B", "x"],
+        ["B", "y"],
+        ["C", "x"],
+        ["C", "y"],
+        ["D", "x"],
+        ["D", "y"]
+      ];
+
+      expect(Array.from(iterable)).toEqual(result);
+    });
+  });
 });

src/index.ts

@@ -60,13 +60,11 @@ export function contains<T>(iterable: Iterable<T>, needle: T) {
 /**
  * Returns an iterable of enumeration pairs.
  */
-export function* enumerate<T>(
+export function enumerate<T>(
   iterable: Iterable<T>,
   offset = 0
 ): Iterable<[number, T]> {
-  let index = offset;
-
-  for (const value of iterable) yield [index++, value];
+  return zip(range(offset), iterable);
 }
 
 /**
@@ -161,8 +159,9 @@ export function* cycle<T>(iterable: Iterable<T>): Iterable<T> {
 /**
  * Make an iterator that repeats `value` over and over again.
  */
-export function* repeat<T>(value: T): Iterable<T> {
-  while (true) yield value;
+export function* repeat<T>(value: T, times?: number): Iterable<T> {
+  if (times === undefined) while (true) yield value;
+  for (let i = 0; i < times; i++) yield value;
 }
 
 /**
@@ -585,3 +584,37 @@ export function max<T extends number>(
 export function sum(iterable: Iterable<number>, start = 0): number {
   return reduce(iterable, (x, y) => x + y, start);
 }
+
+/**
+ * Recursively produce all produces of a list of iterators.
+ */
+function* _product<T>(
+  pools: Iterator<T | typeof SENTINEL>[],
+  buffer: T[] = []
+): Iterable<T[]> {
+  if (pools.length === 0) {
+    yield buffer.slice();
+    return;
+  }
+
+  const [pool, ...others] = pools;
+
+  while (true) {
+    const item = pool.next();
+    if (item.value === SENTINEL) break;
+    buffer.push(item.value);
+    yield* _product(others, buffer);
+    buffer.pop();
+  }
+}
+
+/**
+ * Cartesian product of input iterables.
+ */
+export function* product<T extends any[]>(
+  ...iterables: TupleIterable<T>
+): Iterable<T> {
+  const pools = iterables.map(x => iter(cycle(chain(x, repeat(SENTINEL, 1)))));
+
+  yield* _product(pools) as Iterable<T>;
+}