Repository: pmndrs/zustand
Pull request #466 (2021-06-27 17:03:16)

feat(persist): add a merge function option
Follow up of #459 

This pr adds a new `merge` option for `persist` which can be used to do custom hydration merges.

codesandbox[bot]: This pull request is automatically built and testable in [CodeSandbox](https://codesandbox.io).

  To see build info of the built libraries, click [here](https://ci.codesandbox.io/status/pmndrs/zustand/pr/466/builds/144398) or the icon next to each commit SHA.

Latest deployment of this branch, based on commit ba1d129b466c2a308e9f6456ddacc0911f982a90:

|Sandbox| Source |
|--|--|
|[React](https://codesandbox.io/s/react-1mpwv)| Configuration |
|[React Typescript](https://codesandbox.io/s/react-typescript-sz272)| Configuration |

dai-shi: This may affect #465, so merging now.

Patch:
diff --git a/src/middleware.ts b/src/middleware.ts
index 3cd5d21..6979364 100644
--- a/src/middleware.ts
+++ b/src/middleware.ts
@@ -210,6 +210,11 @@ type PersistOptions<S> = {
    * This function will be called when persisted state versions mismatch with the one specified here.
    */
   migrate?: (persistedState: any, version: number) => S | Promise<S>
+  /**
+   * A function to perform custom hydration merges when combining the stored state with the current one.
+   * By default, this function does a shallow merge.
+   */
+  merge?: (persistedState: any, currentState: S) => S
 }
 
 interface Thenable<Value> {
@@ -264,6 +269,10 @@ export const persist =
       onRehydrateStorage,
       version = 0,
       migrate,
+      merge = (persistedState: any, currentState: S) => ({
+        ...currentState,
+        ...persistedState,
+      }),
     } = options || {}
 
     let storage: StateStorage | undefined
@@ -371,7 +380,10 @@ export const persist =
         }
       })
       .then(() => {
-        stateFromStorageInSync = { ...configResult, ...stateFromStorageInSync }
+        stateFromStorageInSync = merge(
+          stateFromStorageInSync as S,
+          configResult
+        )
         postRehydrationCallback?.(stateFromStorageInSync, undefined)
       })
       .catch((e: Error) => {


Test patch:
diff --git a/tests/persistAsync.test.tsx b/tests/persistAsync.test.tsx
index 9208516..b1db866 100644
--- a/tests/persistAsync.test.tsx
+++ b/tests/persistAsync.test.tsx
@@ -309,4 +309,50 @@ describe('persist middleware with async configuration', () => {
       undefined
     )
   })
+
+  it('can custom merge the stored state', async () => {
+    const storage = {
+      getItem: () =>
+        JSON.stringify({
+          state: {
+            count: 1,
+            actions: {},
+          },
+          version: 0,
+        }),
+      setItem: () => {},
+    }
+
+    const unstorableMethod = () => {}
+
+    const useStore = create(
+      persist(() => ({ count: 0, actions: { unstorableMethod } }), {
+        name: 'test-storage',
+        getStorage: () => storage,
+        merge: (persistedState, currentState) => {
+          delete persistedState.actions
+
+          return {
+            ...currentState,
+            ...persistedState,
+          }
+        },
+      })
+    )
+
+    function Counter() {
+      const { count } = useStore()
+      return <div>count: {count}</div>
+    }
+
+    const { findByText } = render(<Counter />)
+
+    await findByText('count: 1')
+    expect(useStore.getState()).toEqual({
+      count: 1,
+      actions: {
+        unstorableMethod,
+      },
+    })
+  })
 })
diff --git a/tests/persistSync.test.tsx b/tests/persistSync.test.tsx
index 54a060b..78a98b9 100644
--- a/tests/persistSync.test.tsx
+++ b/tests/persistSync.test.tsx
@@ -239,4 +239,42 @@ describe('persist middleware with sync configuration', () => {
     expect(useStore.getState()).toEqual(expectedState)
     expect(onRehydrateStorageSpy).toBeCalledWith(expectedState, undefined)
   })
+
+  it('can custom merge the stored state', () => {
+    const storage = {
+      getItem: () =>
+        JSON.stringify({
+          state: {
+            count: 1,
+            actions: {},
+          },
+          version: 0,
+        }),
+      setItem: () => {},
+    }
+
+    const unstorableMethod = () => {}
+
+    const useStore = create(
+      persist(() => ({ count: 0, actions: { unstorableMethod } }), {
+        name: 'test-storage',
+        getStorage: () => storage,
+        merge: (persistedState, currentState) => {
+          delete persistedState.actions
+
+          return {
+            ...currentState,
+            ...persistedState,
+          }
+        },
+      })
+    )
+
+    expect(useStore.getState()).toEqual({
+      count: 1,
+      actions: {
+        unstorableMethod,
+      },
+    })
+  })
 })