feat(middleware/persist): new storage option (#1463)

* feat(middleware/persist): new storage option

* fix typo

docs/integrations/persisting-store-data.md

@@ -19,7 +19,7 @@ for more details.
 
 ```ts
 import create from 'zustand'
-import { persist } from 'zustand/middleware'
+import { persist, createJSONStorage } from 'zustand/middleware'
 
 export const useBearStore = create(
   persist(
@@ -29,7 +29,7 @@ export const useBearStore = create(
     }),
     {
       name: 'food-storage', // name of the item in the storage (must be unique)
-      getStorage: () => sessionStorage, // (optional) by default, 'localStorage' is used
+      storage: createJSONStorage(() => sessionStorage), // (optional) by default, 'localStorage' is used
     }
   )
 )
@@ -44,7 +44,7 @@ The given name is going to be the key
 used to store your Zustand state in the storage,
 so it must be unique.
 
-### `getStorage`
+### `storage`
 
 > Type: `() => StateStorage`
 
@@ -54,7 +54,7 @@ The `StateStorage` can be imported with:
 import { StateStorage } from 'zustand/middleware'
 ```
 
-> Default: `() => localStorage`
+> Default: `createJSONStorage(() => localStorage)`
 
 Enables you to use your own storage.
 Simply pass a function that returns the storage you want to use.
@@ -62,6 +62,8 @@ Simply pass a function that returns the storage you want to use.
 Example:
 
 ```ts
+import { persist, createJSONStorage } from 'zustand/middleware'
+
 export const useBoundStore = create(
   persist(
     (set, get) => ({
@@ -69,60 +71,7 @@ export const useBoundStore = create(
     }),
     {
       // ...
-      getStorage: () => AsyncStorage,
-    }
-  )
-)
-```
-
-### `serialize`
-
-> Type: `(state: Object) => string | Promise<string>`
-
-> Default: `(state) => JSON.stringify(state)`
-
-The only way to store an object in a storage is as a string.
-If the default method of serialization doesn't suit your needs,
-pass custom functions for serialization
-and [deserialization](#deserialize) (see below).
-
-For example, if you want to store your state in base64:
-
-```ts
-export const useBoundStore = create(
-  persist(
-    (set, get) => ({
-      // ...
-    }),
-    {
-      // ...
-      serialize: (state) => btoa(JSON.stringify(state)),
-    }
-  )
-)
-```
-
-### `deserialize`
-
-> Type: `(str: string) => Object | Promise<Object>`
-
-> Default: `(str) => JSON.parse(str)`
-
-If you pass a custom [`serialize`](#serialize) function,
-you will most likely need to pass a custom deserialize function as well.
-
-To continue the example above,
-you could deserialize the base64 value using the following:
-
-```ts
-export const useBoundStore = create(
-  persist(
-    (set, get) => ({
-      // ...
-    }),
-    {
-      // ...
-      deserialize: (str) => JSON.parse(atob(str)),
+      storage: createJSONStorage(() => AsyncStorage),
     }
   )
 )

readme.md

@@ -310,7 +310,7 @@ You can persist your store's data using any kind of storage.
 
 ```jsx
 import create from 'zustand'
-import { persist } from 'zustand/middleware'
+import { persist, createJSONStorage } from 'zustand/middleware'
 
 const useFishStore = create(
   persist(
@@ -320,7 +320,7 @@ const useFishStore = create(
     }),
     {
       name: 'food-storage', // unique name
-      getStorage: () => sessionStorage, // (optional) by default, 'localStorage' is used
+      storage: createJSONStorage(() => sessionStorage), // (optional) by default, 'localStorage' is used
     }
   )
 )

src/middleware/persist.ts

@@ -11,10 +11,50 @@ type StorageValue<S> = {
   version?: number
 }
 
+export interface PersistStorage<S> {
+  getItem: (
+    name: string
+  ) => StorageValue<S> | null | Promise<StorageValue<S> | null>
+  setItem: (name: string, value: StorageValue<S>) => void | Promise<void>
+  removeItem: (name: string) => void | Promise<void>
+}
+
+export function createJSONStorage<S>(
+  getStorage: () => StateStorage
+): PersistStorage<S> | undefined {
+  let storage: StateStorage | undefined
+  try {
+    storage = getStorage()
+  } catch (e) {
+    // prevent error if the storage is not defined (e.g. when server side rendering a page)
+    return
+  }
+  const persistStorage: PersistStorage<S> = {
+    getItem: (name) => {
+      const parse = (str: string | null) => {
+        if (str === null) {
+          return null
+        }
+        return JSON.parse(str) as StorageValue<S>
+      }
+      const str = (storage as StateStorage).getItem(name) ?? null
+      if (str instanceof Promise) {
+        return str.then(parse)
+      }
+      return parse(str)
+    },
+    setItem: (name, newValue) =>
+      (storage as StateStorage).setItem(name, JSON.stringify(newValue)),
+    removeItem: (name) => (storage as StateStorage).removeItem(name),
+  }
+  return persistStorage
+}
+
 export interface PersistOptions<S, PersistedState = S> {
   /** Name of the storage (must be unique) */
   name: string
   /**
+   * @deprecated Use `storage` instead.
    * A function returning a storage.
    * The storage must fit `window.localStorage`'s api (or an async version of it).
    * For example the storage could be `AsyncStorage` from React Native.
@@ -23,6 +63,7 @@ export interface PersistOptions<S, PersistedState = S> {
    */
   getStorage?: () => StateStorage
   /**
+   * @deprecated Use `storage` instead.
    * Use a custom serializer.
    * The returned string will be stored in the storage.
    *
@@ -30,6 +71,7 @@ export interface PersistOptions<S, PersistedState = S> {
    */
   serialize?: (state: StorageValue<S>) => string | Promise<string>
   /**
+   * @deprecated Use `storage` instead.
    * Use a custom deserializer.
    * Must return an object matching StorageValue<S>
    *
@@ -39,6 +81,15 @@ export interface PersistOptions<S, PersistedState = S> {
   deserialize?: (
     str: string
   ) => StorageValue<PersistedState> | Promise<StorageValue<PersistedState>>
+  /**
+   * Use a custom persist storage.
+   *
+   * Combining `createJSONStorage` helps creating a persist storage
+   * with JSON.parse and JSON.stringify.
+   *
+   * @default createJSONStorage(() => localStorage)
+   */
+  storage?: PersistStorage<S> | undefined
   /**
    * Filter the persisted value.
    *
@@ -76,7 +127,7 @@ type StorePersist<S, Ps> = {
   persist: {
     setOptions: (options: Partial<PersistOptions<S, Ps>>) => void
     clearStorage: () => void
-    rehydrate: () => Promise<void>
+    rehydrate: () => Promise<void> | void
     hasHydrated: () => boolean
     onHydrate: (fn: PersistListener<S>) => () => void
     onFinishHydration: (fn: PersistListener<S>) => () => void
@@ -123,7 +174,7 @@ const toThenable =
     }
   }
 
-const persistImpl: PersistImpl = (config, baseOptions) => (set, get, api) => {
+const oldImpl: PersistImpl = (config, baseOptions) => (set, get, api) => {
   type S = ReturnType<typeof config>
   let options = {
     getStorage: () => localStorage,
@@ -296,6 +347,169 @@ const persistImpl: PersistImpl = (config, baseOptions) => (set, get, api) => {
   return stateFromStorage || configResult
 }
 
+const newImpl: PersistImpl = (config, baseOptions) => (set, get, api) => {
+  type S = ReturnType<typeof config>
+  let options = {
+    storage: createJSONStorage<S>(() => localStorage),
+    partialize: (state: S) => state,
+    version: 0,
+    merge: (persistedState: unknown, currentState: S) => ({
+      ...currentState,
+      ...(persistedState as object),
+    }),
+    ...baseOptions,
+  }
+
+  let hasHydrated = false
+  const hydrationListeners = new Set<PersistListener<S>>()
+  const finishHydrationListeners = new Set<PersistListener<S>>()
+  let storage = options.storage
+
+  if (!storage) {
+    return config(
+      (...args) => {
+        console.warn(
+          `[zustand persist middleware] Unable to update item '${options.name}', the given storage is currently unavailable.`
+        )
+        set(...args)
+      },
+      get,
+      api
+    )
+  }
+
+  const setItem = (): void | Promise<void> => {
+    const state = options.partialize({ ...get() })
+    return (storage as PersistStorage<S>).setItem(options.name, {
+      state,
+      version: options.version,
+    })
+  }
+
+  const savedSetState = api.setState
+
+  api.setState = (state, replace) => {
+    savedSetState(state, replace)
+    void setItem()
+  }
+
+  const configResult = config(
+    (...args) => {
+      set(...args)
+      void setItem()
+    },
+    get,
+    api
+  )
+
+  // a workaround to solve the issue of not storing rehydrated state in sync storage
+  // the set(state) value would be later overridden with initial state by create()
+  // to avoid this, we merge the state from localStorage into the initial state.
+  let stateFromStorage: S | undefined
+
+  // rehydrate initial state with existing stored state
+  const hydrate = () => {
+    if (!storage) return
+
+    hasHydrated = false
+    hydrationListeners.forEach((cb) => cb(get()))
+
+    const postRehydrationCallback =
+      options.onRehydrateStorage?.(get()) || undefined
+
+    // bind is used to avoid `TypeError: Illegal invocation` error
+    return Promise.resolve(storage.getItem.bind(storage)(options.name))
+      .then((deserializedStorageValue) => {
+        if (deserializedStorageValue) {
+          if (
+            typeof deserializedStorageValue.version === 'number' &&
+            deserializedStorageValue.version !== options.version
+          ) {
+            if (options.migrate) {
+              return options.migrate(
+                deserializedStorageValue.state,
+                deserializedStorageValue.version
+              )
+            }
+            console.error(
+              `State loaded from storage couldn't be migrated since no migrate function was provided`
+            )
+          } else {
+            return deserializedStorageValue.state
+          }
+        }
+      })
+      .then((migratedState) => {
+        stateFromStorage = options.merge(
+          migratedState as S,
+          get() ?? configResult
+        )
+
+        set(stateFromStorage as S, true)
+        return setItem()
+      })
+      .then(() => {
+        postRehydrationCallback?.(stateFromStorage, undefined)
+        hasHydrated = true
+        finishHydrationListeners.forEach((cb) => cb(stateFromStorage as S))
+      })
+      .catch((e: Error) => {
+        postRehydrationCallback?.(undefined, e)
+      })
+  }
+
+  ;(api as StoreApi<S> & StorePersist<S, S>).persist = {
+    setOptions: (newOptions) => {
+      options = {
+        ...options,
+        ...newOptions,
+      }
+
+      if (newOptions.storage) {
+        storage = newOptions.storage
+      }
+    },
+    clearStorage: () => {
+      storage?.removeItem(options.name)
+    },
+    getOptions: () => options,
+    rehydrate: () => hydrate(),
+    hasHydrated: () => hasHydrated,
+    onHydrate: (cb) => {
+      hydrationListeners.add(cb)
+
+      return () => {
+        hydrationListeners.delete(cb)
+      }
+    },
+    onFinishHydration: (cb) => {
+      finishHydrationListeners.add(cb)
+
+      return () => {
+        finishHydrationListeners.delete(cb)
+      }
+    },
+  }
+
+  hydrate()
+
+  return stateFromStorage || configResult
+}
+
+const persistImpl: PersistImpl = (config, baseOptions) => {
+  if (
+    'getStorage' in baseOptions ||
+    'serialize' in baseOptions ||
+    'deserialize' in baseOptions
+  ) {
+    console.warn(
+      '[DEPRECATED] `getStorage`, `serialize` and `deserialize` options are deprecated. Please use `storage` option instead.'
+    )
+    return oldImpl(config, baseOptions)
+  }
+  return newImpl(config, baseOptions)
+}
+
 type Persist = <
   T,
   Mps extends [StoreMutatorIdentifier, unknown][] = [],
@@ -303,7 +517,7 @@ type Persist = <
   U = T
 >(
   initializer: StateCreator<T, [...Mps, ['zustand/persist', unknown]], Mcs>,
-  options?: PersistOptions<T, U>
+  options: PersistOptions<T, U>
 ) => StateCreator<T, Mps, [['zustand/persist', U], ...Mcs]>
 
 declare module '../vanilla' {

tests/middlewareTypes.test.tsx

@@ -196,7 +196,7 @@ describe('counter state spec (single middleware)', () => {
     TestComponent
 
     const _testSubtyping: StoreApi<object> = createVanilla(
-      persist(() => ({ count: 0 }))
+      persist(() => ({ count: 0 }), { name: 'prefix' })
     )
   })
 

tests/types.test.tsx

@@ -212,7 +212,7 @@ it('StateCreator<T, [StoreMutatorIdentfier, unknown][]> is StateCreator<T, []>',
     },
   })
 
-  create<State>()(persist(foo()))
+  create<State>()(persist(foo(), { name: 'prefix' }))
 })
 
 it('StateCreator subtyping', () => {
@@ -228,7 +228,7 @@ it('StateCreator subtyping', () => {
     },
   })
 
-  create<State>()(persist(foo()))
+  create<State>()(persist(foo(), { name: 'prefix' }))
 
   const _testSubtyping: StateCreator<State, [['zustand/persist', unknown]]> =
     {} as StateCreator<State, []>