feat(middleware/persist): improve createJSONStorage for Maps (#1763)

* feat: add support for Maps to persist middleware

* refactor: change createJsonStorage to accept custom reviver/replacer

* refactor: map support code review

src/middleware/persist.ts

@@ -23,8 +23,14 @@ export interface PersistStorage<S> {
   removeItem: (name: string) => void | Promise<void>
 }
 
+type JsonStorageOptions = {
+  reviver?: (key: string, value: unknown) => unknown
+  replacer?: (key: string, value: unknown) => unknown
+}
+
 export function createJSONStorage<S>(
-  getStorage: () => StateStorage
+  getStorage: () => StateStorage,
+  options?: JsonStorageOptions
 ): PersistStorage<S> | undefined {
   let storage: StateStorage | undefined
   try {
@@ -39,7 +45,7 @@ export function createJSONStorage<S>(
         if (str === null) {
           return null
         }
-        return JSON.parse(str) as StorageValue<S>
+        return JSON.parse(str, options?.reviver) as StorageValue<S>
       }
       const str = (storage as StateStorage).getItem(name) ?? null
       if (str instanceof Promise) {
@@ -48,7 +54,10 @@ export function createJSONStorage<S>(
       return parse(str)
     },
     setItem: (name, newValue) =>
-      (storage as StateStorage).setItem(name, JSON.stringify(newValue)),
+      (storage as StateStorage).setItem(
+        name,
+        JSON.stringify(newValue, options?.replacer)
+      ),
     removeItem: (name) => (storage as StateStorage).removeItem(name),
   }
   return persistStorage

tests/persistAsync.test.tsx

@@ -3,6 +3,7 @@ import { afterEach, describe, expect, it, jest } from '@jest/globals'
 import { act, render, waitFor } from '@testing-library/react'
 import { create } from 'zustand'
 import { createJSONStorage, persist } from 'zustand/middleware'
+import { replacer, reviver } from './test-utils'
 
 const createPersistantStore = (initialValue: string | null) => {
   let state = initialValue
@@ -698,4 +699,120 @@ describe('persist middleware with async configuration', () => {
     await findByText('count: 2')
     expect(useBoundStore.getState().count).toEqual(2)
   })
+
+  it('can rehydrate state with custom deserialized Map', async () => {
+    const onRehydrateStorageSpy = jest.fn()
+    const storage = {
+      getItem: async () =>
+        JSON.stringify({
+          state: {
+            map: { type: 'Map', value: [['foo', 'bar']] },
+          },
+        }),
+      setItem: () => {},
+      removeItem: () => {},
+    }
+
+    const useBoundStore = create(
+      persist(
+        () => ({
+          map: new Map(),
+        }),
+        {
+          name: 'test-storage',
+          storage: createJSONStorage(() => storage, { replacer, reviver }),
+          onRehydrateStorage: () => onRehydrateStorageSpy,
+        }
+      )
+    )
+
+    function MapDisplay() {
+      const { map } = useBoundStore()
+      return <div>map: {map.get('foo')}</div>
+    }
+
+    const { findByText } = render(
+      <StrictMode>
+        <MapDisplay />
+      </StrictMode>
+    )
+
+    await findByText('map: bar')
+    expect(onRehydrateStorageSpy).toBeCalledWith(
+      { map: new Map([['foo', 'bar']]) },
+      undefined
+    )
+  })
+
+  it('can persist state with custom serialization of Map', async () => {
+    const { storage, setItemSpy } = createPersistantStore(null)
+    const map = new Map()
+
+    const createStore = () => {
+      const onRehydrateStorageSpy = jest.fn()
+      const useBoundStore = create(
+        persist(() => ({ map }), {
+          name: 'test-storage',
+          storage: createJSONStorage(() => storage, { replacer, reviver }),
+          onRehydrateStorage: () => onRehydrateStorageSpy,
+        })
+      )
+      return { useBoundStore, onRehydrateStorageSpy }
+    }
+
+    // Initialize from empty storage
+    const { useBoundStore, onRehydrateStorageSpy } = createStore()
+
+    function MapDisplay() {
+      const { map } = useBoundStore()
+      return <div>map-content: {map.get('foo')}</div>
+    }
+
+    const { findByText } = render(
+      <StrictMode>
+        <MapDisplay />
+      </StrictMode>
+    )
+    await findByText('map-content:')
+    await waitFor(() => {
+      expect(onRehydrateStorageSpy).toBeCalledWith({ map }, undefined)
+    })
+
+    // Write something to the store
+    const updatedMap = new Map(map).set('foo', 'bar')
+    act(() => useBoundStore.setState({ map: updatedMap }))
+    await findByText('map-content: bar')
+
+    expect(setItemSpy).toBeCalledWith(
+      'test-storage',
+      JSON.stringify({
+        state: { map: { type: 'Map', value: [['foo', 'bar']] } },
+        version: 0,
+      })
+    )
+
+    // Create the same store a second time and check if the persisted state
+    // is loaded correctly
+    const {
+      useBoundStore: useBoundStore2,
+      onRehydrateStorageSpy: onRehydrateStorageSpy2,
+    } = createStore()
+    function MapDisplay2() {
+      const { map } = useBoundStore2()
+      return <div>map-content: {map.get('foo')}</div>
+    }
+
+    const { findByText: findByText2 } = render(
+      <StrictMode>
+        <MapDisplay2 />
+      </StrictMode>
+    )
+    await findByText2('map-content: bar')
+    await waitFor(() => {
+      expect(onRehydrateStorageSpy2).toBeCalledWith(
+        { map: updatedMap },
+        undefined
+      )
+    })
+  })
 })

tests/persistSync.test.tsx

@@ -1,6 +1,7 @@
 import { afterEach, describe, expect, it, jest } from '@jest/globals'
 import { create } from 'zustand'
 import { createJSONStorage, persist } from 'zustand/middleware'
+import { replacer, reviver } from './test-utils'
 
 const createPersistentStore = (initialValue: string | null) => {
   let state = initialValue
@@ -658,4 +659,84 @@ describe('persist middleware with sync configuration', () => {
 
     expect(useBoundStore.getState().count).toEqual(2)
   })
+
+  it('can rehydrate state with custom deserialized Map', () => {
+    const storage = {
+      getItem: () =>
+        JSON.stringify({
+          map: { type: 'Map', value: [['foo', 'bar']] },
+        }),
+      setItem: () => {},
+      removeItem: () => {},
+    }
+
+    const map = new Map()
+    const onRehydrateStorageSpy = jest.fn()
+    const useBoundStore = create(
+      persist(
+        () => ({
+          map,
+        }),
+        {
+          name: 'test-storage',
+          storage: createJSONStorage(() => storage),
+          onRehydrateStorage: () => onRehydrateStorageSpy,
+        }
+      )
+    )
+
+    const updatedMap = map.set('foo', 'bar')
+    expect(useBoundStore.getState()).toEqual({
+      map: updatedMap,
+    })
+    expect(onRehydrateStorageSpy).toBeCalledWith({ map: updatedMap }, undefined)
+  })
+
+  it('can persist state with custom serialization of Map', () => {
+    const { storage, setItemSpy } = createPersistentStore(null)
+    const map = new Map()
+
+    const createStore = () => {
+      const onRehydrateStorageSpy = jest.fn()
+      const useBoundStore = create(
+        persist(() => ({ map }), {
+          name: 'test-storage',
+          storage: createJSONStorage(() => storage, { replacer, reviver }),
+          onRehydrateStorage: () => onRehydrateStorageSpy,
+        })
+      )
+      return { useBoundStore, onRehydrateStorageSpy }
+    }
+
+    // Initialize from empty storage
+    const { useBoundStore, onRehydrateStorageSpy } = createStore()
+    expect(useBoundStore.getState()).toEqual({ map })
+    expect(onRehydrateStorageSpy).toBeCalledWith({ map }, undefined)
+
+    // Write something to the store
+    const updatedMap = map.set('foo', 'bar')
+    useBoundStore.setState({ map: updatedMap })
+    expect(useBoundStore.getState()).toEqual({
+      map: updatedMap,
+    })
+    expect(setItemSpy).toBeCalledWith(
+      'test-storage',
+      JSON.stringify({
+        state: { map: { type: 'Map', value: [['foo', 'bar']] } },
+        version: 0,
+      })
+    )
+
+    // Create the same store a second time and check if the persisted state
+    // is loaded correctly
+    const {
+      useBoundStore: useBoundStore2,
+      onRehydrateStorageSpy: onRehydrateStorageSpy2,
+    } = createStore()
+    expect(useBoundStore2.getState()).toEqual({ map: updatedMap })
+    expect(onRehydrateStorageSpy2).toBeCalledWith(
+      { map: updatedMap },
+      undefined
+    )
+  })
 })

tests/test-utils.ts

@@ -0,0 +1,33 @@
+type ReplacedMap = {
+  type: 'Map'
+  value: [string, unknown][]
+}
+
+export const replacer = (
+  key: string,
+  value: unknown
+): ReplacedMap | unknown => {
+  if (value instanceof Map) {
+    return {
+      type: 'Map',
+      value: Array.from(value.entries()),
+    }
+  } else {
+    return value
+  }
+}
+
+export const reviver = (key: string, value: ReplacedMap | unknown) => {
+  if (isReplacedMap(value)) {
+    return new Map(value.value)
+  }
+  return value
+}
+
+const isReplacedMap = (value: any): value is ReplacedMap => {
+  if (value && value.type === 'Map') {
+    return true
+  }
+
+  return false
+}