From 3cc667fff2f09e6e8333c25ddad7847a14ae5a81 Mon Sep 17 00:00:00 2001
From: hw <xodltus@naver.com>
Date: Mon, 14 Oct 2024 21:50:01 +0900
Subject: [PATCH] feat: support Promise object to`fx`

---
 src/Lazy/fx.ts             | 58 ++++++++++++++++++++++++++++++++++----
 test/Lazy/fx.spec.ts       | 30 +++++++++++++++++++-
 type-check/Lazy/fx.test.ts |  5 ++++
 3 files changed, 87 insertions(+), 6 deletions(-)

diff --git a/src/Lazy/fx.ts b/src/Lazy/fx.ts
index 6265c128..2ad96760 100644
--- a/src/Lazy/fx.ts
+++ b/src/Lazy/fx.ts
@@ -1,4 +1,4 @@
-import { isAsyncIterable, isIterable } from "../_internal/utils";
+import { isAsyncIterable, isIterable, isPromise } from "../_internal/utils";
 import consume from "../consume";
 import each from "../each";
 import every from "../every";
@@ -593,6 +593,8 @@ export class FxIterable<A> {
   }
 }
 
+type NonPromise<T> = T extends Promise<unknown> ? never : T;
+
 /**
  * `fx` allows functions provided by existing `fxts` to be used in a method chaining.
  *  Not all functions are provided as methods and can be connected through `chain` if necessary.
@@ -621,18 +623,64 @@ export class FxIterable<A> {
  *   .toArray(); // [11, 12, 13, 14]
  * ```
  */
-function fx<T extends Iterable<unknown> | AsyncIterable<unknown>>(
+function fx<
+  T extends
+    | Promise<Iterable<unknown>>
+    | Iterable<unknown>
+    | AsyncIterable<unknown>,
+>(
   a: T,
 ): T extends Iterable<unknown>
   ? FxIterable<IterableInfer<T>>
-  : FxAsyncIterable<IterableInfer<T>> {
+  : T extends Promise<Iterable<infer U>>
+  ? FxAsyncIterable<U>
+  : FxAsyncIterable<IterableInfer<NonPromise<T>>> {
   if (isAsyncIterable(a)) {
     return new FxAsyncIterable(a) as any;
   } else if (isIterable(a)) {
     return new FxIterable(a) as any;
+  } else if (isPromise(a)) {
+    let started = false;
+    let itered = false;
+    let iter: any;
+
+    const next = async (): Promise<{
+      done: boolean;
+      value: any;
+    }> => {
+      if (!started) {
+        iter = await a;
+        started = true;
+      }
+
+      if (itered) {
+        return iter.next();
+      }
+
+      if (isIterable(iter)) {
+        iter = iter[Symbol.iterator]();
+        itered = true;
+        return next();
+      } else {
+        return { done: true, value: iter };
+      }
+    };
+
+    const asyncIter = {
+      [Symbol.asyncIterator]() {
+        return this;
+      },
+      async next() {
+        return next();
+      },
+    };
+
+    return new FxAsyncIterable(asyncIter) as any;
+  } else {
+    throw new TypeError(
+      `'fx' must be type of Iterable or AsyncIterable or Promise`,
+    );
   }
-
-  throw new TypeError(`'fx' must be type of Iterable or AsyncIterable`);
 }
 
 export default fx;
diff --git a/test/Lazy/fx.spec.ts b/test/Lazy/fx.spec.ts
index ef817d7e..58b5da70 100644
--- a/test/Lazy/fx.spec.ts
+++ b/test/Lazy/fx.spec.ts
@@ -1,4 +1,4 @@
-import { fx, size, toArray, toAsync, uniq } from "../../src";
+import { delay, fx, size, toArray, toAsync, uniq } from "../../src";
 
 describe("fx", function () {
   describe("sync", function () {
@@ -58,5 +58,33 @@ describe("fx", function () {
 
       expect(arrSize).toEqual(3);
     });
+
+    it("handle Promise object", async function () {
+      const res = await fx(Promise.resolve([1, 2, 3, 4])).toArray();
+      expect(res).toEqual([1, 2, 3, 4]);
+    });
+
+    it("handle Promise.all object", async function () {
+      const res1 = fx(
+        Promise.all([
+          Promise.resolve(1),
+          Promise.resolve(2),
+          Promise.resolve(3),
+          Promise.resolve(4),
+        ]),
+      );
+      expect(await res1.toArray()).toEqual([1, 2, 3, 4]);
+
+      const res2 = await fx(
+        Promise.all([
+          delay(1000, 1),
+          Promise.resolve(2),
+          delay(1000, 3),
+          Promise.resolve(4),
+        ]),
+      ).toArray();
+
+      expect(res2).toEqual([1, 2, 3, 4]);
+    }, 1050);
   });
 });
diff --git a/type-check/Lazy/fx.test.ts b/type-check/Lazy/fx.test.ts
index b572c1dc..3164389c 100644
--- a/type-check/Lazy/fx.test.ts
+++ b/type-check/Lazy/fx.test.ts
@@ -24,6 +24,9 @@ const res7 = fx(toAsync([1, 2, 3]))
 const res8 = [...fx([1, 2, 3])];
 const res9 = [...fx("abc")];
 
+const res10 = fx(Promise.resolve([1]));
+const res11 = fx(Promise.all([Promise.resolve(1)]));
+
 checks([
   check<typeof res1, Cast<Iterable<number>, typeof res1>, Test.Pass>(),
   check<typeof res2, number[], Test.Pass>(),
@@ -34,4 +37,6 @@ checks([
   check<typeof res7, Promise<number[]>, Test.Pass>(),
   check<typeof res8, number[], Test.Pass>(),
   check<typeof res9, string[], Test.Pass>(),
+  check<typeof res10, Cast<AsyncIterable<number>, typeof res10>, Test.Pass>(),
+  check<typeof res11, Cast<AsyncIterable<number>, typeof res11>, Test.Pass>(),
 ]);