feat: transfer start-unit declaration to initialProps

src/index.ts

@@ -5,7 +5,8 @@ const { createDomain, createStore, createEffect, createEvent } = domain;
 export * from "effector";
 
 export { domain, createDomain, createStore, createEffect, createEvent };
+export * from "./lib/constants";
+export * from "./types/custom";
 export * from "./with-hydrate";
 export * from "./with-start";
 export * from "./with-fork";
-export * from "./types";

src/lib/get-start-units.ts

@@ -1,30 +1,28 @@
-import { DocumentContext } from "next/document";
-import { NextComponentType } from "next";
-import { Unit } from "effector";
+import { AppInitialProps } from "next/app";
+import { Unit, is } from "effector";
 
-import { PageContext } from "../types";
+import { AppType, Enhancer, RenderPage, PageContext } from "../types";
 
 import { START_UNIT_KEY } from "./constants";
 
-/* eslint-disable @typescript-eslint/no-explicit-any */
-
 type StartUnits = Array<Unit<PageContext>>;
-type RenderPage = DocumentContext["renderPage"];
-type Enhancer<C extends NextComponentType = NextComponentType<any, any, any>> = (Component: C) => C;
 
 export function getStartUnits(originalRenderPage: RenderPage) {
   const units: StartUnits = [];
 
   originalRenderPage({ enhanceApp: getStartUnit(units) });
-  originalRenderPage({ enhanceComponent: getStartUnit(units) });
 
-  return units;
+  return units.filter(is.unit);
 }
 
-function getStartUnit(units: StartUnits): Enhancer {
-  return (Component) => () => {
-    if (START_UNIT_KEY in Component) {
-      units.push(Component[START_UNIT_KEY]);
+function getStartUnit<P extends AppInitialProps>(units: StartUnits): Enhancer<AppType<P>> {
+  return () => (props) => {
+    if (START_UNIT_KEY in props) {
+      units.push(props[START_UNIT_KEY]);
+    }
+
+    if (START_UNIT_KEY in props.pageProps) {
+      units.push(props.pageProps[START_UNIT_KEY]);
     }
 
     return null;

src/lib/render-page-with-scope.tsx

@@ -1,13 +1,9 @@
 import * as React from "react";
-import { AppContext, AppInitialProps } from "next/app";
-import { DocumentContext } from "next/document";
 import { Provider } from "effector-react/ssr";
-import { NextComponentType } from "next";
+import { AppInitialProps } from "next/app";
 import { Scope } from "effector/fork";
 
-type Enhancer<C> = (Component: C) => C;
-type RenderPage = DocumentContext["renderPage"];
-type AppType<P extends AppInitialProps> = NextComponentType<AppContext, AppInitialProps, P>;
+import { AppType, Enhancer, RenderPage } from "../types";
 
 export function renderPageWithScope(scope: Scope, originalRenderPage: RenderPage): RenderPage {
   return (params) => {

src/types/extracted.ts

@@ -0,0 +1,7 @@
+import { AppContext, AppInitialProps } from "next/app";
+import { DocumentContext } from "next/document";
+import { NextComponentType } from "next";
+
+export type Enhancer<C> = (Component: C) => C;
+export type RenderPage = DocumentContext["renderPage"];
+export type AppType<P extends AppInitialProps> = NextComponentType<AppContext, AppInitialProps, P>;

src/types/index.ts

@@ -0,0 +1,2 @@
+export * from "./extracted";
+export * from "./custom";

src/with-fork.tsx

@@ -14,7 +14,7 @@ type InitialStateKey = typeof INITIAL_STATE_KEY;
 type InitialState = ReturnType<typeof serialize>;
 type ExtendedNextData = NextData & { [key in InitialStateKey]: InitialState };
 
-interface WrappedDocumentProps extends DocumentProps {
+interface CustomDocumentProps extends DocumentProps {
   initialState: InitialState;
   __NEXT_DATA__: ExtendedNextData;
 }
@@ -25,7 +25,7 @@ export interface WithForkConfig {
 
 export function withFork({ debug }: WithForkConfig = {}) {
   return (Document: typeof NextDocument) =>
-    class WithForkDocument extends React.Component<WrappedDocumentProps> {
+    class WithForkDocument extends React.Component<CustomDocumentProps> {
       static renderDocument = Document.renderDocument;
       static headTagsMiddleware = Document.headTagsMiddleware;
       static bodyTagsMiddleware = Document.bodyTagsMiddleware;
@@ -80,7 +80,7 @@ export function withFork({ debug }: WithForkConfig = {}) {
         };
       }
 
-      constructor(props: WrappedDocumentProps) {
+      constructor(props: CustomDocumentProps) {
         super(props);
 
         props.__NEXT_DATA__[INITIAL_STATE_KEY] = props.initialState;

src/with-start.ts

@@ -21,11 +21,9 @@ export function withStart(unit: Unit<PageContext>) {
         initialProps = await originalGetInitialProps(ctx);
       }
 
-      return initialProps;
+      return Object.assign({}, initialProps, { [START_UNIT_KEY]: unit });
     };
 
-    component[START_UNIT_KEY] = unit;
-
     return component;
   };
 }