implemented wrapping react lazy components

XantreDev · Oct 6, 2023 · f057209 · f057209
1 parent a35a602
commit f057209
Show file tree

Hide file tree

Showing 6 changed files with 107 additions and 17 deletions.
diff --git a/.changeset/quiet-sloths-care.md b/.changeset/quiet-sloths-care.md
@@ -0,0 +1,5 @@
+---
+"react-fast-hoc": minor
+---
+
+Added support for components wrapped with `React.lazy`
diff --git a/packages/react-fast-hoc/package.json b/packages/react-fast-hoc/package.json
@@ -28,6 +28,7 @@
     "build": "pnpm build:collect-readme & rollup -c"
   },
   "devDependencies": {
+    "@rollup/plugin-typescript": "^11.1.1",
     "@testing-library/react": "^14.0.0",
     "@types/node": "^20.2.5",
     "@types/react": "^18.2.0",
@@ -37,17 +38,17 @@
     "hotscript": "^1.0.12",
     "jsdom": "^21.1.1",
     "nanobundle": "^1.6.0",
+    "radash": "^11.0.0",
     "react": "^18.2.0",
     "react-dom": "^18.2.0",
-    "ts-toolbelt": "^9.6.0",
-    "tsconfig": "workspace:*",
-    "type-fest": "^3.9.0",
-    "vitest": "^0.30.1",
-    "@rollup/plugin-typescript": "^11.1.1",
     "rollup": "^3.23.1",
     "rollup-plugin-dts": "^5.3.0",
     "rollup-plugin-esbuild": "^5.0.0",
-    "rollup-plugin-node-externals": "^6.1.1"
+    "rollup-plugin-node-externals": "^6.1.1",
+    "ts-toolbelt": "^9.6.0",
+    "tsconfig": "workspace:*",
+    "type-fest": "^3.9.0",
+    "vitest": "^0.30.1"
   },
   "peerDependencies": {
     "hotscript": "^1.0.12",

diff --git a/packages/react-fast-hoc/src/handlers/mimicToNewComponent.ts b/packages/react-fast-hoc/src/handlers/mimicToNewComponent.ts
@@ -1,4 +1,4 @@
-export class MimicToNewComponentHandler implements ProxyHandler<Function> {
+export class MimicToNewComponentHandler implements ProxyHandler<object> {
   private _componentProps = new WeakMap<Function, Map<PropertyKey, unknown>>();
 
   get(target: Function, p: PropertyKey, receiver: any) {

diff --git a/packages/react-fast-hoc/src/index.test.ts b/packages/react-fast-hoc/src/index.test.ts
@@ -1,4 +1,5 @@
 import { cleanup, render, waitFor } from "@testing-library/react";
+import { sleep } from "radash";
 import { Objects } from "hotscript";
 import React, { ComponentType, createElement, forwardRef, memo } from "react";
 import { Function } from "ts-toolbelt";
@@ -64,7 +65,6 @@ describe("transformProps", () => {
 
   class ClassComponent extends React.Component {
     constructor(props: unknown) {
-      console.log("class created", props);
       super(props as {});
       propsDetector(this.props);
     }
@@ -140,16 +140,69 @@ describe("transformProps", () => {
       const Cmp = vi.fn(Component);
       const Lazy = React.lazy(() => Promise.resolve({ default: Cmp }));
 
-      console.log(Lazy._payload?._result?.toString());
-      console.log(Lazy._init?.toString());
       render(
         createElement(React.Suspense, {}, createElement(addBebeHoc(Lazy)))
       );
       await waitFor(() => {
         expect(Cmp).toHaveBeenCalled();
-        console.log(Lazy);
         expect(addBebeProp).toHaveBeenCalled();
-        expect(addBebeProp).lastCalledWith({});
+        expect(addBebeProp).lastReturnedWith({
+          bebe: true,
+        });
+      });
+    });
+    test("pending lazy", async () => {
+      const Cmp = vi.fn(Component);
+      const lazyInit = vi.fn(() => sleep(20).then(() => ({ default: Cmp })));
+      const Lazy = React.lazy(lazyInit);
+      const r = render(createElement(React.Suspense, {}, createElement(Lazy)));
+      expect(Cmp).not.toHaveBeenCalled();
+      expect(lazyInit).toHaveBeenCalled();
+
+      r.rerender(
+        createElement(React.Suspense, {}, createElement(addBebeHoc(Lazy)))
+      );
+
+      await waitFor(
+        () => {
+          expect(Cmp).toHaveBeenCalled();
+          expect(addBebeProp).toHaveBeenCalled();
+          expect(addBebeProp).lastReturnedWith({
+            bebe: true,
+          });
+        },
+        {
+          timeout: 100,
+        }
+      );
+    });
+    test("resolved lazy", async () => {
+      const Cmp = vi.fn(Component);
+      const Lazy = React.lazy(() => Promise.resolve({ default: Cmp }));
+      const r = render(createElement(React.Suspense, {}, createElement(Lazy)));
+      await waitFor(
+        () => {
+          expect(Cmp).toHaveBeenCalled();
+          expect(addBebeProp).not.toHaveBeenCalled();
+        },
+        {
+          timeout: 100,
+        }
+      );
+
+      r.rerender(
+        createElement(React.Suspense, {}, createElement(addBebeHoc(Lazy)))
+      );
+      expect(Cmp).toHaveBeenCalledTimes(2);
+      expect(addBebeProp).toHaveBeenCalled();
+      expect(Cmp).lastCalledWith(
+        {
+          bebe: true,
+        },
+        {}
+      );
+      expect(addBebeProp).lastReturnedWith({
+        bebe: true,
       });
     });
   });

diff --git a/packages/react-fast-hoc/src/internals.ts b/packages/react-fast-hoc/src/internals.ts
@@ -54,8 +54,12 @@ type RealComponentType<TProps extends object, IRef = unknown> =
     }
   | {
       $$typeof: typeof REACT_LAZY_TYPE;
-      _status: -1 | 0 | 1 | 2;
-      _result: unknown;
+      _payload: {
+        _status: -1 | 0 | 1 | 2;
+        _result: unknown;
+      };
+      // returns component or throws promise
+      _init: (arg: unknown) => React.ComponentType<unknown>;
     }
   | React.ComponentClass<TProps>
   | React.FC<TProps>;
@@ -99,7 +103,7 @@ export const wrapComponentIntoHoc = <TProps extends object>(
   Component: RealComponentType<TProps>,
   handler: HocTransformer,
   mimicToNewComponentHandler: null | MimicToNewComponentHandler
-) => {
+): unknown => {
   // this case assumes that it's ClassComponent
   if (isClassComponent(Component)) {
     return wrapFunctionalFROrDefault(
@@ -128,11 +132,30 @@ export const wrapComponentIntoHoc = <TProps extends object>(
     };
   }
   if ("$$typeof" in Component && Component["$$typeof"] === REACT_LAZY_TYPE) {
-    return Component;
+    let result: RealComponentType<any>;
+    return {
+      $$typeof: REACT_LAZY_TYPE,
+      _payload: Component._payload,
+      _init: (arg: unknown) => {
+        const initRes = Component._init(arg);
+        if (!result) {
+          result = wrapComponentIntoHoc(
+            initRes,
+            handler,
+            mimicToNewComponentHandler
+          ) as RealComponentType<any>;
+        }
+        return result;
+      },
+    } as RealComponentType<any>;
   }
 
   const proxied = new Proxy(Component, handler);
+
   return mimicToNewComponentHandler
-    ? new Proxy(proxied, mimicToNewComponentHandler)
+    ? (new Proxy(
+        proxied,
+        mimicToNewComponentHandler
+      ) as RealComponentType<TProps>)
     : proxied;
 };
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml