Repository: pmndrs/zustand (ts)

Issue:
Persist middleware api
Follow up of https://github.com/pmndrs/zustand/pull/248#issuecomment-752130586. The idea is to make an api to interact with `persist`.

#### Change the options
```js
useMyStore.persist.setOptions({
  getStorage: () => sessionStorage,
})
```
This `setOptions` could be useful to, for example, choose a storage at runtime if the use check a box or something. It could also be used to change the name of the storage as needed here: https://github.com/pmndrs/zustand/pull/248#issuecomment-752130586

#### Clear storage data
```js
useMyStore.persist.clearStorage()
```

#### On demand rehydration (see #374)
```js
useMyStore.persist.rehydrate()
```

Pull request description:
feat(middleware): add persist api
Closes #267, #513

Patch:
diff --git a/src/index.ts b/src/index.ts
index c13e566..61fb247 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -1,4 +1,10 @@
-import { useEffect, useLayoutEffect, useReducer, useRef } from 'react'
+import {
+  useDebugValue,
+  useEffect,
+  useLayoutEffect,
+  useReducer,
+  useRef,
+} from 'react'
 import createImpl, {
   EqualityChecker,
   GetState,
@@ -127,9 +133,11 @@ export default function create<
       return unsubscribe
     }, [])
 
-    return hasNewStateSlice
+    const sliceToReturn = hasNewStateSlice
       ? (newStateSlice as StateSlice)
       : currentSliceRef.current
+    useDebugValue(sliceToReturn)
+    return sliceToReturn
   }
 
   Object.assign(useStore, api)
diff --git a/src/middleware.ts b/src/middleware.ts
index 9cceb40..24cb229 100644
--- a/src/middleware.ts
+++ b/src/middleware.ts
@@ -4,7 +4,6 @@ import {
   PartialState,
   SetState,
   State,
-  StateCreator,
   StateListener,
   StateSelector,
   StateSliceListener,
@@ -287,6 +286,8 @@ type DeepPartial<T extends Object> = {
 export type StateStorage = {
   getItem: (name: string) => string | null | Promise<string | null>
   setItem: (name: string, value: string) => void | Promise<void>
+  // Note: This will be required in v4
+  removeItem?: (name: string) => void | Promise<void>
 }
 type StorageValue<S> = { state: DeepPartial<S>; version?: number }
 export type PersistOptions<
@@ -360,6 +361,15 @@ export type PersistOptions<
    */
   merge?: (persistedState: any, currentState: S) => S
 }
+export type StoreApiWithPersist<S extends State> = StoreApi<S> & {
+  persist: {
+    setOptions: (options: Partial<PersistOptions<S>>) => void
+    clearStorage: () => void
+    rehydrate: () => Promise<void>
+    hasHydrated: () => boolean
+    onHydrate: (fn: (state: S) => void) => void
+  }
+}
 
 interface Thenable<Value> {
   then<V>(
@@ -401,72 +411,88 @@ const toThenable =
   }
 
 export const persist =
-  <S extends State>(config: StateCreator<S>, options: PersistOptions<S>) =>
-  (set: SetState<S>, get: GetState<S>, api: StoreApi<S>): S => {
-    const {
-      name,
-      getStorage = () => localStorage,
-      serialize = JSON.stringify as (state: StorageValue<S>) => string,
-      deserialize = JSON.parse as (str: string) => StorageValue<Partial<S>>,
-      blacklist,
-      whitelist,
-      partialize = (state: S) => state,
-      onRehydrateStorage,
-      version = 0,
-      migrate,
-      merge = (persistedState: any, currentState: S) => ({
+  <
+    S extends State,
+    CustomSetState extends SetState<S>,
+    CustomGetState extends GetState<S>,
+    CustomStoreApi extends StoreApiWithPersist<S>
+  >(
+    config: (
+      set: CustomSetState,
+      get: CustomGetState,
+      api: CustomStoreApi
+    ) => S,
+    baseOptions: PersistOptions<S>
+  ) =>
+  (set: CustomSetState, get: CustomGetState, api: CustomStoreApi): S => {
+    let options = {
+      getStorage: () => localStorage,
+      serialize: JSON.stringify as (state: StorageValue<S>) => string,
+      deserialize: JSON.parse as (str: string) => StorageValue<Partial<S>>,
+      partialize: (state: S) => state,
+      version: 0,
+      merge: (persistedState: any, currentState: S) => ({
         ...currentState,
         ...persistedState,
       }),
-    } = options || {}
+      ...baseOptions,
+    }
 
-    if (blacklist || whitelist) {
+    if (options.blacklist || options.whitelist) {
       console.warn(
         `The ${
-          blacklist ? 'blacklist' : 'whitelist'
+          options.blacklist ? 'blacklist' : 'whitelist'
         } option is deprecated and will be removed in the next version. Please use the 'partialize' option instead.`
       )
     }
 
+    let hasHydrated = false
+    const onHydrateCallbacks: Parameters<
+      StoreApiWithPersist<S>['persist']['onHydrate']
+    >[0][] = []
     let storage: StateStorage | undefined
 
     try {
-      storage = getStorage()
+      storage = options.getStorage()
     } catch (e) {
       // prevent error if the storage is not defined (e.g. when server side rendering a page)
     }
 
     if (!storage) {
       return config(
-        (...args) => {
+        ((...args) => {
           console.warn(
-            `Persist middleware: unable to update ${name}, the given storage is currently unavailable.`
+            `[zustand persist middleware] Unable to update item '${options.name}', the given storage is currently unavailable.`
           )
           set(...args)
-        },
+        }) as CustomSetState,
         get,
         api
       )
+    } else if (!storage.removeItem) {
+      console.warn(
+        `[zustand persist middleware] The given storage for item '${options.name}' does not contain a 'removeItem' method, which will be required in v4.`
+      )
     }
 
-    const thenableSerialize = toThenable(serialize)
+    const thenableSerialize = toThenable(options.serialize)
 
     const setItem = (): Thenable<void> => {
-      const state = partialize({ ...get() })
+      const state = options.partialize({ ...get() })
 
-      if (whitelist) {
+      if (options.whitelist) {
         ;(Object.keys(state) as (keyof S)[]).forEach((key) => {
-          !whitelist.includes(key) && delete state[key]
+          !options.whitelist?.includes(key) && delete state[key]
         })
       }
-      if (blacklist) {
-        blacklist.forEach((key) => delete state[key])
+      if (options.blacklist) {
+        options.blacklist.forEach((key) => delete state[key])
       }
 
       let errorInSync: Error | undefined
-      const thenable = thenableSerialize({ state, version })
+      const thenable = thenableSerialize({ state, version: options.version })
         .then((serializedValue) =>
-          (storage as StateStorage).setItem(name, serializedValue)
+          (storage as StateStorage).setItem(options.name, serializedValue)
         )
         .catch((e) => {
           errorInSync = e
@@ -485,60 +511,93 @@ export const persist =
     }
 
     const configResult = config(
-      (...args) => {
+      ((...args) => {
         set(...args)
         void setItem()
-      },
+      }) as CustomSetState,
       get,
       api
     )
 
-    // rehydrate initial state with existing stored state
-
     // 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
-    const postRehydrationCallback = onRehydrateStorage?.(get()) || undefined
-    // bind is used to avoid `TypeError: Illegal invocation` error
-    toThenable(storage.getItem.bind(storage))(name)
-      .then((storageValue) => {
-        if (storageValue) {
-          return deserialize(storageValue)
-        }
-      })
-      .then((deserializedStorageValue) => {
-        if (deserializedStorageValue) {
-          if (
-            typeof deserializedStorageValue.version === 'number' &&
-            deserializedStorageValue.version !== version
-          ) {
-            if (migrate) {
-              return migrate(
-                deserializedStorageValue.state,
-                deserializedStorageValue.version
+
+    // rehydrate initial state with existing stored state
+    const hydrate = () => {
+      if (!storage) return
+
+      hasHydrated = false
+
+      const postRehydrationCallback =
+        options.onRehydrateStorage?.(get()) || undefined
+
+      // bind is used to avoid `TypeError: Illegal invocation` error
+      return toThenable(storage.getItem.bind(storage))(options.name)
+        .then((storageValue) => {
+          if (storageValue) {
+            return options.deserialize(storageValue)
+          }
+        })
+        .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
             }
-            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, configResult)
+
+          set(stateFromStorage as S, true)
+          return setItem()
+        })
+        .then(() => {
+          postRehydrationCallback?.(stateFromStorage, undefined)
+          hasHydrated = true
+          onHydrateCallbacks.forEach((cb) => cb(stateFromStorage as S))
+        })
+        .catch((e: Error) => {
+          postRehydrationCallback?.(undefined, e)
+        })
+    }
+
+    api.persist = {
+      setOptions: (newOptions) => {
+        options = {
+          ...options,
+          ...newOptions,
         }
-      })
-      .then((migratedState) => {
-        stateFromStorage = merge(migratedState as S, configResult)
 
-        set(stateFromStorage as S, true)
-        return setItem()
-      })
-      .then(() => {
-        postRehydrationCallback?.(stateFromStorage, undefined)
-      })
-      .catch((e: Error) => {
-        postRehydrationCallback?.(undefined, e)
-      })
+        if (newOptions.getStorage) {
+          storage = newOptions.getStorage()
+        }
+      },
+      clearStorage: () => {
+        storage?.removeItem?.(options.name)
+      },
+      rehydrate: () => hydrate() as Promise<void>,
+      hasHydrated: () => hasHydrated,
+      onHydrate: (cb) => {
+        onHydrateCallbacks.push(cb)
+      },
+    }
+
+    hydrate()
 
     return stateFromStorage || configResult
   }


Test patch:
diff --git a/tests/middlewareTypes.test.tsx b/tests/middlewareTypes.test.tsx
index ce4e44e..d5dfacb 100644
--- a/tests/middlewareTypes.test.tsx
+++ b/tests/middlewareTypes.test.tsx
@@ -1,7 +1,14 @@
 import { produce } from 'immer'
 import type { Draft } from 'immer'
-import create, { GetState, SetState, State, StateCreator } from 'zustand'
+import create, {
+  GetState,
+  SetState,
+  State,
+  StateCreator,
+  StoreApi,
+} from 'zustand'
 import {
+  StoreApiWithPersist,
   StoreApiWithSubscribeWithSelector,
   combine,
   devtools,
@@ -17,7 +24,14 @@ type TImmerConfigFn<T extends State> = (
 type TImmerConfig<T extends State> = StateCreator<T, TImmerConfigFn<T>>
 
 const immer =
-  <T extends State>(config: TImmerConfig<T>): StateCreator<T> =>
+  <
+    T extends State,
+    CustomSetState extends SetState<T>,
+    CustomGetState extends GetState<T>,
+    CustomStoreApi extends StoreApi<T>
+  >(
+    config: TImmerConfig<T>
+  ): StateCreator<T, CustomSetState, CustomGetState, CustomStoreApi> =>
   (set, get, api) =>
     config(
       (partial, replace) => {
@@ -174,7 +188,12 @@ describe('counter state spec (single middleware)', () => {
   })
 
   it('persist', () => {
-    const useStore = create<CounterState>(
+    const useStore = create<
+      CounterState,
+      SetState<CounterState>,
+      GetState<CounterState>,
+      StoreApiWithPersist<CounterState>
+    >(
       persist(
         (set, get) => ({
           count: 1,
@@ -326,7 +345,12 @@ describe('counter state spec (double middleware)', () => {
   })
 
   it('devtools & persist', () => {
-    const useStore = create<CounterState>(
+    const useStore = create<
+      CounterState,
+      SetState<CounterState>,
+      GetState<CounterState>,
+      StoreApiWithPersist<CounterState>
+    >(
       devtools(
         persist(
           (set, get) => ({
@@ -354,7 +378,12 @@ describe('counter state spec (double middleware)', () => {
 
 describe('counter state spec (triple middleware)', () => {
   it('devtools & persist & immer', () => {
-    const useStore = create<CounterState>(
+    const useStore = create<
+      CounterState,
+      SetState<CounterState>,
+      GetState<CounterState>,
+      StoreApiWithPersist<CounterState>
+    >(
       devtools(
         persist(
           immer((set, get) => ({
@@ -413,7 +442,8 @@ describe('counter state spec (triple middleware)', () => {
       CounterState,
       SetState<CounterState>,
       GetState<CounterState>,
-      StoreApiWithSubscribeWithSelector<CounterState>
+      StoreApiWithSubscribeWithSelector<CounterState> &
+        StoreApiWithPersist<CounterState>
     >(
       devtools(
         subscribeWithSelector(
@@ -451,7 +481,8 @@ describe('counter state spec (quadruple middleware)', () => {
       CounterState,
       SetState<CounterState>,
       GetState<CounterState>,
-      StoreApiWithSubscribeWithSelector<CounterState>
+      StoreApiWithSubscribeWithSelector<CounterState> &
+        StoreApiWithPersist<CounterState>
     >(
       devtools(
         subscribeWithSelector(
@@ -485,3 +516,33 @@ describe('counter state spec (quadruple middleware)', () => {
     TestComponent
   })
 })
+
+describe('more complex state spec with subscribeWithSelector', () => {
+  it('#619', () => {
+    type MyState = {
+      foo: boolean | string
+    }
+    const useStore = create<
+      MyState,
+      SetState<MyState>,
+      GetState<MyState>,
+      StoreApiWithSubscribeWithSelector<MyState>
+    >(
+      subscribeWithSelector(() => ({
+        // Note: It complains without type assertion.
+        foo: true as boolean | string,
+      }))
+    )
+    const TestComponent = () => {
+      useStore((s) => s.foo)
+      useStore().foo
+      useStore.getState().foo
+      useStore.subscribe(
+        (state) => state.foo,
+        (foo) => console.log(foo)
+      )
+      return <></>
+    }
+    TestComponent
+  })
+})
diff --git a/tests/persistAsync.test.tsx b/tests/persistAsync.test.tsx
index 03c9c86..dba5c52 100644
--- a/tests/persistAsync.test.tsx
+++ b/tests/persistAsync.test.tsx
@@ -14,13 +14,20 @@ const createPersistantStore = (initialValue: string | null) => {
     state = newState
   }
 
+  const removeItem = async (name: string) => {
+    removeItemSpy(name)
+    state = null
+  }
+
   const getItemSpy = jest.fn()
   const setItemSpy = jest.fn()
+  const removeItemSpy = jest.fn()
 
   return {
-    storage: { getItem, setItem },
+    storage: { getItem, setItem, removeItem },
     getItemSpy,
     setItemSpy,
+    removeItemSpy,
   }
 }
 
@@ -39,6 +46,7 @@ describe('persist middleware with async configuration', () => {
           version: 0,
         }),
       setItem: () => {},
+      removeItem: () => {},
     }
 
     const useStore = create(
@@ -82,6 +90,7 @@ describe('persist middleware with async configuration', () => {
         throw new Error('getItem error')
       },
       setItem: () => {},
+      removeItem: () => {},
     }
 
     const useStore = create(
@@ -175,6 +184,7 @@ describe('persist middleware with async configuration', () => {
           version: 12,
         }),
       setItem: setItemSpy,
+      removeItem: () => {},
     }
 
     const useStore = create(
@@ -217,6 +227,7 @@ describe('persist middleware with async configuration', () => {
           version: 12,
         }),
       setItem: (_: string, _value: string) => {},
+      removeItem: () => {},
     }
 
     const useStore = create(
@@ -253,6 +264,7 @@ describe('persist middleware with async configuration', () => {
           version: 12,
         }),
       setItem: () => {},
+      removeItem: () => {},
     }
 
     const useStore = create(
@@ -293,6 +305,7 @@ describe('persist middleware with async configuration', () => {
           version: 0,
         }),
       setItem: () => {},
+      removeItem: () => {},
     }
 
     const unstorableMethod = () => {}
@@ -332,6 +345,7 @@ describe('persist middleware with async configuration', () => {
           version: 0,
         }),
       setItem: () => {},
+      removeItem: () => {},
     }
 
     const unstorableMethod = () => {}
@@ -376,6 +390,7 @@ describe('persist middleware with async configuration', () => {
           },
         }),
       setItem: () => {},
+      removeItem: () => {},
     }
 
     const useStore = create(
@@ -398,4 +413,48 @@ describe('persist middleware with async configuration', () => {
       count: 1,
     })
   })
+
+  it('can manually rehydrate through the api', async () => {
+    const storageValue = '{"state":{"count":1},"version":0}'
+
+    const storage = {
+      getItem: async () => '',
+      setItem: () => {},
+      removeItem: () => {},
+    }
+
+    const useStore = create(
+      persist(() => ({ count: 0 }), {
+        name: 'test-storage',
+        getStorage: () => storage,
+      })
+    )
+
+    storage.getItem = async () => storageValue
+    await useStore.persist.rehydrate()
+    expect(useStore.getState()).toEqual({
+      count: 1,
+    })
+  })
+
+  it('can check if the store has been hydrated through the api', async () => {
+    const storage = {
+      getItem: async () => null,
+      setItem: () => {},
+      removeItem: () => {},
+    }
+
+    const useStore = create(
+      persist(() => ({ count: 0 }), {
+        name: 'test-storage',
+        getStorage: () => storage,
+      })
+    )
+    expect(useStore.persist.hasHydrated()).toBe(false)
+    await new Promise((resolve) => useStore.persist.onHydrate(resolve))
+    expect(useStore.persist.hasHydrated()).toBe(true)
+
+    await useStore.persist.rehydrate()
+    expect(useStore.persist.hasHydrated()).toBe(true)
+  })
 })
diff --git a/tests/persistSync.test.tsx b/tests/persistSync.test.tsx
index d4d46d3..beb2248 100644
--- a/tests/persistSync.test.tsx
+++ b/tests/persistSync.test.tsx
@@ -12,12 +12,17 @@ const createPersistantStore = (initialValue: string | null) => {
     setItemSpy(name, newState)
     state = newState
   }
+  const removeItem = (name: string) => {
+    removeItemSpy(name)
+    state = null
+  }
 
   const getItemSpy = jest.fn()
   const setItemSpy = jest.fn()
+  const removeItemSpy = jest.fn()
 
   return {
-    storage: { getItem, setItem },
+    storage: { getItem, setItem, removeItem },
     getItemSpy,
     setItemSpy,
   }
@@ -37,6 +42,7 @@ describe('persist middleware with sync configuration', () => {
           version: 0,
         }),
       setItem: () => {},
+      removeItem: () => {},
     }
 
     const onRehydrateStorageSpy = jest.fn()
@@ -67,6 +73,7 @@ describe('persist middleware with sync configuration', () => {
         throw new Error('getItem error')
       },
       setItem: () => {},
+      removeItem: () => {},
     }
 
     const spy = jest.fn()
@@ -131,6 +138,7 @@ describe('persist middleware with sync configuration', () => {
           version: 12,
         }),
       setItem: setItemSpy,
+      removeItem: () => {},
     }
 
     const useStore = create(
@@ -165,6 +173,7 @@ describe('persist middleware with sync configuration', () => {
           version: 12,
         }),
       setItem: (_: string, _value: string) => {},
+      removeItem: () => {},
     }
 
     const useStore = create(
@@ -191,6 +200,7 @@ describe('persist middleware with sync configuration', () => {
           version: 12,
         }),
       setItem: () => {},
+      removeItem: () => {},
     }
 
     const useStore = create(
@@ -222,6 +232,7 @@ describe('persist middleware with sync configuration', () => {
           version: 0,
         }),
       setItem: () => {},
+      removeItem: () => {},
     }
 
     const unstorableMethod = () => {}
@@ -251,6 +262,7 @@ describe('persist middleware with sync configuration', () => {
           version: 0,
         }),
       setItem: () => {},
+      removeItem: () => {},
     }
 
     const unstorableMethod = () => {}
@@ -287,6 +299,7 @@ describe('persist middleware with sync configuration', () => {
           },
         }),
       setItem: () => {},
+      removeItem: () => {},
     }
 
     const useStore = create(
@@ -308,6 +321,7 @@ describe('persist middleware with sync configuration', () => {
     const storage = {
       getItem: () => '',
       setItem: setItemSpy,
+      removeItem: () => {},
     }
 
     const useStore = create(
@@ -365,4 +379,131 @@ describe('persist middleware with sync configuration', () => {
       })
     )
   })
+
+  it('can change the options through the api', () => {
+    const setItemSpy = jest.fn()
+
+    const storage = {
+      getItem: () => null,
+      setItem: setItemSpy,
+      removeItem: () => {},
+    }
+
+    const useStore = create(
+      persist(() => ({ count: 0 }), {
+        name: 'test-storage',
+        getStorage: () => storage,
+      })
+    )
+
+    useStore.setState({})
+    expect(setItemSpy).toBeCalledWith(
+      'test-storage',
+      '{"state":{"count":0},"version":0}'
+    )
+
+    useStore.persist.setOptions({
+      name: 'test-storage-2',
+      partialize: (state) =>
+        Object.fromEntries(
+          Object.entries(state).filter(([key]) => key !== 'count')
+        ),
+    })
+    useStore.setState({})
+    expect(setItemSpy).toBeCalledWith(
+      'test-storage-2',
+      '{"state":{},"version":0}'
+    )
+  })
+
+  it('can clear the storage through the api', () => {
+    const removeItemSpy = jest.fn()
+
+    const storage = {
+      getItem: () => null,
+      setItem: () => {},
+      removeItem: removeItemSpy,
+    }
+
+    const useStore = create(
+      persist(() => ({ count: 0 }), {
+        name: 'test-storage',
+        getStorage: () => storage,
+      })
+    )
+
+    useStore.persist.clearStorage()
+    expect(removeItemSpy).toBeCalledWith('test-storage')
+  })
+
+  it('can manually rehydrate through the api', () => {
+    const storageValue = '{"state":{"count":1},"version":0}'
+
+    const storage = {
+      getItem: () => '',
+      setItem: () => {},
+      removeItem: () => {},
+    }
+
+    const useStore = create(
+      persist(() => ({ count: 0 }), {
+        name: 'test-storage',
+        getStorage: () => storage,
+      })
+    )
+
+    storage.getItem = () => storageValue
+    useStore.persist.rehydrate()
+    expect(useStore.getState()).toEqual({
+      count: 1,
+    })
+  })
+
+  it('can check if the store has been hydrated through the api', async () => {
+    const storage = {
+      getItem: () => null,
+      setItem: () => {},
+      removeItem: () => {},
+    }
+
+    const useStore = create(
+      persist(() => ({ count: 0 }), {
+        name: 'test-storage',
+        getStorage: () => storage,
+      })
+    )
+
+    expect(useStore.persist.hasHydrated()).toBe(true)
+
+    await useStore.persist.rehydrate()
+    expect(useStore.persist.hasHydrated()).toBe(true)
+  })
+
+  it('can wait for hydration through the api', async () => {
+    const storageValue = '{"state":{"count":1},"version":0}'
+
+    const onHydrateSpy1 = jest.fn()
+    const onHydrateSpy2 = jest.fn()
+
+    const storage = {
+      getItem: () => '',
+      setItem: () => {},
+      removeItem: () => {},
+    }
+
+    const useStore = create(
+      persist(() => ({ count: 0 }), {
+        name: 'test-storage',
+        getStorage: () => storage,
+      })
+    )
+
+    useStore.persist.onHydrate(onHydrateSpy1)
+    useStore.persist.onHydrate(onHydrateSpy2)
+
+    storage.getItem = () => storageValue
+    await useStore.persist.rehydrate()
+    expect(onHydrateSpy1).toBeCalledWith({ count: 1 })
+    expect(onHydrateSpy2).toBeCalledWith({ count: 1 })
+  })
 })