support persist middleware

dai-shi · dai-shi · commit fa334f94f38f · 2024-07-13T11:31:08.000+09:00
diff --git a/README.md b/README.md
@@ -16,20 +16,23 @@ npm install zustand zustand-valtio valtio
 ## Usage
 
 ```jsx
-import { createWithProxy } from 'zustand-valtio';
-
-const useCounterState = createWithProxy({
-  count: 0,
-  inc() {
-    this.count++;
-  },
-});
+import { create } from 'zustand';
+import { withProxy } from 'zustand-valtio';
+
+const useCounterState = create(
+  withProxy({
+    count: 0,
+    inc() {
+      this.count++;
+    },
+  }),
+);
 
 const Counter = () => {
   const count = useCounterState((state) => state.count);
   const inc = useCounterState((state) => state.inc);
   // Or this works too
-  // const inc = () => ++useCounterState.proxyState.count;
+  // const inc = () => ++useCounterState.getProxyState().count;
   return (
     <>
       <div>Count: {count}</div>
@@ -60,6 +63,7 @@ and open <http://localhost:8080> in your web browser.
 You can also try them directly:
 [01](https://stackblitz.com/github/zustandjs/zustand-valtio/tree/main/examples/01_counter)
 [02](https://stackblitz.com/github/zustandjs/zustand-valtio/tree/main/examples/02_methods)
+[03](https://stackblitz.com/github/zustandjs/zustand-valtio/tree/main/examples/03_middleware)
 
 ## Tweets
 
diff --git a/examples/03_middleware/index.html b/examples/03_middleware/index.html
@@ -0,0 +1,9 @@
+<html>
+  <head>
+    <title>example</title>
+  </head>
+  <body>
+    <div id="root"></div>
+    <script type="module" src="/src/main.tsx"></script>
+  </body>
+</html>
diff --git a/examples/03_middleware/package.json b/examples/03_middleware/package.json
@@ -0,0 +1,22 @@
+{
+  "name": "example",
+  "version": "0.0.0",
+  "private": true,
+  "type": "module",
+  "dependencies": {
+    "react": "latest",
+    "react-dom": "latest",
+    "valtio": "latest",
+    "zustand": "latest",
+    "zustand-valtio": "latest"
+  },
+  "devDependencies": {
+    "@types/react": "latest",
+    "@types/react-dom": "latest",
+    "typescript": "latest",
+    "vite": "latest"
+  },
+  "scripts": {
+    "dev": "vite"
+  }
+}
diff --git a/examples/03_middleware/src/app.tsx b/examples/03_middleware/src/app.tsx
@@ -0,0 +1,36 @@
+import { create } from 'zustand';
+import { persist } from 'zustand/middleware';
+import { withProxy } from 'zustand-valtio';
+
+const useCounterState = create(
+  persist(
+    withProxy({
+      count: 0,
+      inc() {
+        this.count++;
+      },
+    }),
+    { name: 'counter' },
+  ),
+);
+
+const Counter = () => {
+  const count = useCounterState((state) => state.count);
+  const inc = useCounterState((state) => state.inc);
+  return (
+    <>
+      <div>Count: {count}</div>
+      <button type="button" onClick={inc}>
+        +1
+      </button>
+    </>
+  );
+};
+
+const App = () => (
+  <div>
+    <Counter />
+  </div>
+);
+
+export default App;
diff --git a/examples/03_middleware/src/main.tsx b/examples/03_middleware/src/main.tsx
@@ -0,0 +1,10 @@
+import { StrictMode } from 'react';
+import { createRoot } from 'react-dom/client';
+
+import App from './app';
+
+createRoot(document.getElementById('root')!).render(
+  <StrictMode>
+    <App />
+  </StrictMode>,
+);
diff --git a/examples/03_middleware/tsconfig.json b/examples/03_middleware/tsconfig.json
@@ -0,0 +1,14 @@
+{
+  "compilerOptions": {
+    "strict": true,
+    "target": "es2018",
+    "esModuleInterop": true,
+    "module": "esnext",
+    "moduleResolution": "bundler",
+    "skipLibCheck": true,
+    "allowJs": true,
+    "noUncheckedIndexedAccess": true,
+    "exactOptionalPropertyTypes": true,
+    "jsx": "react-jsx"
+  }
+}
diff --git a/package.json b/package.json
@@ -41,7 +41,8 @@
     "test:types:examples": "tsc -p examples --noEmit",
     "test:spec": "vitest run",
     "examples:01_counter": "DIR=01_counter vite",
-    "examples:02_methods": "DIR=02_methods vite"
+    "examples:02_methods": "DIR=02_methods vite",
+    "examples:03_middleware": "DIR=03_middleware vite"
   },
   "keywords": [
     "react",
diff --git a/src/index.ts b/src/index.ts
@@ -37,34 +37,51 @@ type WithProxyImpl = <T extends object>(
   initialObject: T,
 ) => StateCreator<Snapshot<T>, [], []>;
 
-const withProxyImpl: WithProxyImpl = (initialObject) => (set, _get, api) => {
+const withProxyImpl: WithProxyImpl = (initialObject) => (set, get, api) => {
   type AnyObject = Record<string, unknown>;
   const proxyState = proxy(initialObject);
   let mutating = 0;
+  let updating = 0;
   const updateState = () => {
     if (!mutating) {
+      ++updating;
       set(snapshot(proxyState), true);
+      --updating;
     }
   };
   Object.keys(proxyState).forEach((key) => {
     const fn = (proxyState as AnyObject)[key];
+    // TODO this doesn't handle nested objects
     if (typeof fn === 'function') {
       (proxyState as AnyObject)[key] = (...args: never[]) => {
         try {
-          mutating += 1;
+          ++mutating;
           return fn.apply(proxyState, args);
         } finally {
-          mutating -= 1;
+          --mutating;
           updateState();
         }
       };
     }
   });
-  type Api = StoreApi<Snapshot<typeof initialObject>> &
-    StoreWithProxy<typeof initialObject>;
+  type Api = StoreApi<unknown> & StoreWithProxy<typeof initialObject>;
   delete (api as Api).setState;
   (api as Api).getProxyState = () => proxyState;
   subscribe(proxyState, updateState, true);
+  api.subscribe(() => {
+    if (!updating) {
+      // HACK for persist hydration
+      const state = get() as AnyObject;
+      Object.keys(state).forEach((key) => {
+        const val = state[key];
+        // TODO this doesn't handle nested objects
+        if (typeof val !== 'function') {
+          // XXX this will throw if val is a snapshot
+          (proxyState as AnyObject)[key] = val;
+        }
+      });
+    }
+  });
   return snapshot(proxyState);
 };
 
