fix: flatMap should accept any (async)iterable

Author: LumaKernel

README.md

@@ -3,6 +3,51 @@
 Strict and wrapper version implementation for
 https://github.com/tc39/proposal-iterator-helpers.
 
+## Usage
+
+```ts
+import {
+  asyncIteratorFrom,
+  iteratorFrom,
+  wrapAsyncIterator,
+  wrapIterator,
+} from "./mod.ts";
+
+function* naturals() {
+  let i = 0;
+  while (true) {
+    yield i;
+    i += 1;
+  }
+}
+
+const arr1 = wrapIterator(naturals())
+  .filter((n) => n % 2 === 1) // filter odd numbers
+  .map((n) => n ** 2) // square numbers
+  .flatMap((n) => [n, n]) // twice each numbers
+  .take(10) // cut up to 10 items
+  .toArray(); // evaluate and collect items into array
+console.log(arr1); // [1, 1, 9, 9, 25, 25, 49, 49, 81, 81]
+
+const arr2 = await wrapAsyncIterator(asyncIteratorFrom(naturals()))
+  .filter(async (n) => {
+    const res = await fetch(`https://api.isevenapi.xyz/api/iseven/${n}/`);
+    if (!res.body) throw new Error("No body");
+    const raw = Uint8Array.from(
+      await wrapAsyncIterator(asyncIteratorFrom(res.body))
+        .flatMap((e) => e)
+        .toArray(),
+    );
+    const obj = JSON.parse(new TextDecoder().decode(raw));
+    return obj.iseven;
+  }) // filter even numbers
+  .map((n) => n ** 2) // square numbers
+  .flatMap((n) => [n, n]) // twice each numbers
+  .take(10) // cut up to 10 items
+  .toArray(); // evaluate and collect items into array
+console.log(arr2); // [0, 0, 4, 4, 16, 16, 36, 36, 64, 64]
+```
+
 ## Goals
 
 - Implement all proposed features with wrapper API.

lib/wrap_async_iterator.ts

@@ -15,7 +15,13 @@ export type WrappedAsyncIterator<T> = {
   drop: (limit: number) => WrappedAsyncIterator<T>;
   asIndexedPairs: () => WrappedAsyncIterator<[number, T]>;
   flatMap: <U>(
-    mapperFn: (t: T) => U | U[] | Promise<U | U[]>,
+    mapperFn: (
+      t: T,
+    ) =>
+      | U
+      | Iterable<U>
+      | AsyncIterable<U>
+      | Promise<U | Iterable<U> | AsyncIterable<U>>,
   ) => WrappedAsyncIterator<U>;
   reduce: {
     <U>(
@@ -114,11 +120,25 @@ export const wrapAsyncIterator = <T>(
       }
       return wrapAsyncIterator((async function* () {
         for await (const v of { [Symbol.asyncIterator]: () => ite }) {
-          const arr = await mapperFn(v);
-          if (Array.isArray(arr)) {
-            for (const e of arr) yield e;
+          const inner = await mapperFn(v);
+          const getAsync = (inner as any)[Symbol.asyncIterator];
+          if (getAsync != null) {
+            if (typeof getAsync !== "function") {
+              throw new TypeError(`${getAsync} is not a function`);
+            }
+            const iteInner = getAsync.call(inner);
+            for await (const vInner of iteInner) yield vInner;
           } else {
-            yield arr;
+            const getSync = (inner as any)[Symbol.iterator];
+            if (getSync != null) {
+              if (typeof getSync !== "function") {
+                throw new TypeError(`${getSync} is not a function`);
+              }
+              const iteInner = getSync.call(inner);
+              for (const vInner of iteInner) yield vInner;
+            } else {
+              yield inner;
+            }
           }
         }
       })());

lib/wrap_iterator.ts

@@ -15,7 +15,7 @@ export type WrappedIterator<T> = {
   drop: (limit: number) => WrappedIterator<T>;
   asIndexedPairs: () => WrappedIterator<[number, T]>;
   flatMap: <U>(
-    mapperFn: (t: T) => U | U[],
+    mapperFn: (t: T) => U | Iterable<U>,
   ) => WrappedIterator<U>;
   reduce: {
     <U>(
@@ -112,11 +112,16 @@ export const wrapIterator = <T>(ite: Iterator<T>): WrappedIterator<T> => {
       }
       return wrapIterator((function* () {
         for (const v of { [Symbol.iterator]: () => ite }) {
-          const arr = mapperFn(v);
-          if (Array.isArray(arr)) {
-            for (const e of arr) yield e;
+          const inner = mapperFn(v);
+          const getSync = (inner as any)[Symbol.iterator];
+          if (getSync != null) {
+            if (typeof getSync !== "function") {
+              throw new TypeError(`${getSync} is not a function`);
+            }
+            const iteInner = getSync.call(inner);
+            for (const vInner of iteInner) yield vInner;
           } else {
-            yield arr;
+            yield inner;
           }
         }
       })());

mod.ts

@@ -1,4 +1,4 @@
-export * from "./lib/async_iterator_from.spec.ts";
-export * from "./lib/iterator_from.spec.ts";
-export * from "./lib/wrap_async_iterator.spec.ts";
-export * from "./lib/wrap_iterator.spec.ts";
+export * from "./lib/async_iterator_from.ts";
+export * from "./lib/iterator_from.ts";
+export * from "./lib/wrap_async_iterator.ts";
+export * from "./lib/wrap_iterator.ts";

test/unit/wrap_async_iterator.spec.ts

@@ -213,10 +213,17 @@ Deno.test({
   name: "AsyncIterator.prototype.flatMap",
   async fn() {
     asserts.assertEquals(
-      await wrapAsyncIterator(asyncIteratorFrom([1, 2, [3, 4], [5, [6, 7], 8]]))
+      await wrapAsyncIterator(asyncIteratorFrom([
+        1,
+        2,
+        [3, 4],
+        [5, [6, 7], 8],
+        // any (async)iterable should be accepted
+        Uint8Array.from([9, 10, 11]),
+      ]))
         .flatMap<number | number[]>((e) => typeof e === "number" ? e * 100 : e)
         .toArray(),
-      [100, 200, 3, 4, 5, [6, 7], 8],
+      [100, 200, 3, 4, 5, [6, 7], 8, 9, 10, 11],
     );
     asserts.assertThrows(
       () => wrapAsyncIterator(asyncIteratorFrom([1, 2, 3])).flatMap(1 as any),

test/unit/wrap_iterator.spec.ts

@@ -211,10 +211,17 @@ Deno.test({
   name: "Iterator.prototype.flatMap",
   fn() {
     asserts.assertEquals(
-      wrapIterator(iteratorFrom([1, 2, [3, 4], [5, [6, 7], 8]]))
+      wrapIterator(iteratorFrom([
+        1,
+        2,
+        [3, 4],
+        [5, [6, 7], 8],
+        // any iterable should be accepted
+        Uint8Array.from([9, 10, 11]),
+      ]))
         .flatMap<number | number[]>((e) => typeof e === "number" ? e * 100 : e)
         .toArray(),
-      [100, 200, 3, 4, 5, [6, 7], 8],
+      [100, 200, 3, 4, 5, [6, 7], 8, 9, 10, 11],
     );
     asserts.assertThrows(
       () => wrapIterator(iteratorFrom([1, 2, 3])).flatMap(1 as any),