Repository: pmndrs/zustand
Pull request #646 (2021-11-07 17:53:58)

feat(middleware): add onFinishHydration to persist api
I noticed in #644 that I kind of misnamed `useStore.persist.onHydrate` :grimacing: 
By the name `onHydrate` we would most likely expect the callback to be called when hydration starts, while in reality it's called when hydration is finished.
I fixed this by actually calling `onHydrate`'s callback when hydration starts, and adding a new `onFinishHydration` method to the persist api.

This pr also adds a new `useStore.persist.useHydration()` hook to make it easier to (mainly?) know if hydration is finished.

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/646/builds/186102) or the icon next to each commit SHA.

Latest deployment of this branch, based on commit e1334e5bd1d4d5c43ad1b089d0aeb27b711412d7:

|Sandbox| Source |
|--|--|
|[React](https://codesandbox.io/s/react-qwh5s)| Configuration |
|[React Typescript](https://codesandbox.io/s/react-typescript-8oln9)| Configuration |
|[React Browserify](https://codesandbox.io/s/react-browserify-n8h5g)| Configuration |
|[React Snowpack](https://codesandbox.io/s/react-snowpack-51ndi)| Configuration |
|[React Parcel](https://codesandbox.io/s/react-parcel-tkej2)| Configuration |
|[Next.js](https://codesandbox.io/s/next-js-gjy1j)| Configuration |

github-actions[bot]: **Size Change:** +287 B (+3%) 

**Total Size:** 11.4 kB

| Filename | Size | Change |  |
| :--- | :---: | :---: | :---: |
| `dist/esm/middleware.js` | 2.91 kB | +136 B (+5%) | 🔍 |
| `dist/middleware.js` | 3.08 kB | +151 B (+5%) | 🔍 |

<details><summary>ℹ️ <strong>View Unchanged</strong></summary>

| Filename | Size |
| :--- | :---: |
| `dist/context.js` | 622 B |
| `dist/esm/context.js` | 521 B |
| `dist/esm/index.js` | 1.23 kB |
| `dist/esm/shallow.js` | 272 B |
| `dist/esm/vanilla.js` | 520 B |
| `dist/index.js` | 1.35 kB |
| `dist/shallow.js` | 318 B |
| `dist/vanilla.js` | 630 B |

</details>



<a href="https://github.com/preactjs/compressed-size-action"><sub>compressed-size-action</sub></a>

Patch:
diff --git a/readme.md b/readme.md
index 670d5e5..bb27724 100644
--- a/readme.md
+++ b/readme.md
@@ -213,6 +213,30 @@ const unsub4 = useStore.subscribe(state => [state.paw, state.fur], console.log,
 const unsub5 = useStore.subscribe(state => state.paw, console.log, { fireImmediately: true })
 ```
 
+<details>
+<summary>How to type store with `subscribeWithSelector` in TypeScript</summary>
+
+```ts
+import create, { GetState, SetState } from 'zustand'
+import { StoreApiWithSubscribeWithSelector } from 'zustand/middleware'
+
+type BearState = {
+  paw: boolean
+  snout: boolean
+  fur: boolean
+}
+const useStore = create<
+  BearState,
+  SetState<BearState>,
+  GetState<BearState>,
+  StoreApiWithSubscribeWithSelector<BearState>
+>(subscribeWithSelector(() => ({ paw: true, snout: true, fur: true })))
+```
+
+For more complex typing with multiple middlewares,
+Please refer [middlewareTypes.test.tsx](./tests/middlewareTypes.test.tsx).
+</details>
+
 ## Using zustand without React
 
 Zustands core can be imported and used without the React dependency. The only difference is that the create function does not return a hook, but the api utilities.
diff --git a/src/index.ts b/src/index.ts
index 61fb247..9f17bee 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -44,11 +44,28 @@ export type UseBoundStore<
   <U>(selector: StateSelector<T, U>, equalityFn?: EqualityChecker<U>): U
 } & CustomStoreApi
 
-export default function create<
+function create<
   TState extends State,
-  CustomSetState = SetState<TState>,
-  CustomGetState = GetState<TState>,
-  CustomStoreApi extends StoreApi<TState> = StoreApi<TState>
+  CustomSetState,
+  CustomGetState,
+  CustomStoreApi extends StoreApi<TState>
+>(
+  createState:
+    | StateCreator<TState, CustomSetState, CustomGetState, CustomStoreApi>
+    | CustomStoreApi
+): UseBoundStore<TState, CustomStoreApi>
+
+function create<TState extends State>(
+  createState:
+    | StateCreator<TState, SetState<TState>, GetState<TState>, any>
+    | StoreApi<TState>
+): UseBoundStore<TState, StoreApi<TState>>
+
+function create<
+  TState extends State,
+  CustomSetState,
+  CustomGetState,
+  CustomStoreApi extends StoreApi<TState>
 >(
   createState:
     | StateCreator<TState, CustomSetState, CustomGetState, CustomStoreApi>
@@ -158,3 +175,5 @@ export default function create<
 
   return useStore
 }
+
+export default create
diff --git a/src/middleware.ts b/src/middleware.ts
index 218844a..49deb3f 100644
--- a/src/middleware.ts
+++ b/src/middleware.ts
@@ -367,13 +367,15 @@ export type PersistOptions<
    */
   merge?: (persistedState: any, currentState: S) => S
 }
+type PersistListener<S> = (state: S) => void
 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
+    onHydrate: (fn: PersistListener<S>) => () => void
+    onFinishHydration: (fn: PersistListener<S>) => () => void
   }
 }
 
@@ -457,9 +459,8 @@ export const persist =
     }
 
     let hasHydrated = false
-    const onHydrateCallbacks: Parameters<
-      StoreApiWithPersist<S>['persist']['onHydrate']
-    >[0][] = []
+    const hydrationListeners = new Set<PersistListener<S>>()
+    const finishHydrationListeners = new Set<PersistListener<S>>()
     let storage: StateStorage | undefined
 
     try {
@@ -539,6 +540,7 @@ export const persist =
       if (!storage) return
 
       hasHydrated = false
+      hydrationListeners.forEach((cb) => cb(get()))
 
       const postRehydrationCallback =
         options.onRehydrateStorage?.(get()) || undefined
@@ -579,7 +581,7 @@ export const persist =
         .then(() => {
           postRehydrationCallback?.(stateFromStorage, undefined)
           hasHydrated = true
-          onHydrateCallbacks.forEach((cb) => cb(stateFromStorage as S))
+          finishHydrationListeners.forEach((cb) => cb(stateFromStorage as S))
         })
         .catch((e: Error) => {
           postRehydrationCallback?.(undefined, e)
@@ -603,7 +605,18 @@ export const persist =
       rehydrate: () => hydrate() as Promise<void>,
       hasHydrated: () => hasHydrated,
       onHydrate: (cb) => {
-        onHydrateCallbacks.push(cb)
+        hydrationListeners.add(cb)
+
+        return () => {
+          hydrationListeners.delete(cb)
+        }
+      },
+      onFinishHydration: (cb) => {
+        finishHydrationListeners.add(cb)
+
+        return () => {
+          finishHydrationListeners.delete(cb)
+        }
       },
     }
 
diff --git a/src/vanilla.ts b/src/vanilla.ts
index 96657fe..b1e38ff 100644
--- a/src/vanilla.ts
+++ b/src/vanilla.ts
@@ -52,11 +52,29 @@ export type StateCreator<
   CustomStoreApi extends StoreApi<T> = StoreApi<T>
 > = (set: CustomSetState, get: CustomGetState, api: CustomStoreApi) => T
 
-export default function create<
+function create<
   TState extends State,
-  CustomSetState = SetState<TState>,
-  CustomGetState = GetState<TState>,
-  CustomStoreApi extends StoreApi<TState> = StoreApi<TState>
+  CustomSetState,
+  CustomGetState,
+  CustomStoreApi extends StoreApi<TState>
+>(
+  createState: StateCreator<
+    TState,
+    CustomSetState,
+    CustomGetState,
+    CustomStoreApi
+  >
+): CustomStoreApi
+
+function create<TState extends State>(
+  createState: StateCreator<TState, SetState<TState>, GetState<TState>, any>
+): StoreApi<TState>
+
+function create<
+  TState extends State,
+  CustomSetState,
+  CustomGetState,
+  CustomStoreApi extends StoreApi<TState>
 >(
   createState: StateCreator<
     TState,
@@ -131,3 +149,5 @@ export default function create<
   )
   return api as unknown as CustomStoreApi
 }
+
+export default create


Test patch:
diff --git a/tests/middlewareTypes.test.tsx b/tests/middlewareTypes.test.tsx
index b76d364..38c7c01 100644
--- a/tests/middlewareTypes.test.tsx
+++ b/tests/middlewareTypes.test.tsx
@@ -205,6 +205,29 @@ describe('counter state spec (single middleware)', () => {
         { name: 'prefix' }
       )
     )
+    const TestComponent = () => {
+      useStore((s) => s.count) * 2
+      useStore((s) => s.inc)()
+      useStore().count * 2
+      useStore().inc()
+      useStore.getState().count * 2
+      useStore.getState().inc()
+      useStore.persist.hasHydrated()
+      return <></>
+    }
+    TestComponent
+  })
+
+  it('persist without custom api (#638)', () => {
+    const useStore = create<CounterState>(
+      persist(
+        (set, get) => ({
+          count: 1,
+          inc: () => set({ count: get().count + 1 }, false),
+        }),
+        { name: 'prefix' }
+      )
+    )
     const TestComponent = () => {
       useStore((s) => s.count) * 2
       useStore((s) => s.inc)()
@@ -383,6 +406,7 @@ describe('counter state spec (double middleware)', () => {
       useStore.getState().count * 2
       useStore.getState().inc()
       useStore.setState({ count: 0 }, false, 'reset')
+      useStore.persist.hasHydrated()
       return <></>
     }
     TestComponent
@@ -419,6 +443,7 @@ describe('counter state spec (triple middleware)', () => {
       useStore.getState().count * 2
       useStore.getState().inc()
       useStore.setState({ count: 0 }, false, 'reset')
+      useStore.persist.hasHydrated()
       return <></>
     }
     TestComponent
@@ -486,6 +511,7 @@ describe('counter state spec (triple middleware)', () => {
         (count) => console.log(count * 2)
       )
       useStore.setState({ count: 0 }, false, 'reset')
+      useStore.persist.hasHydrated()
       return <></>
     }
     TestComponent
@@ -530,6 +556,7 @@ describe('counter state spec (quadruple middleware)', () => {
         (count) => console.log(count * 2)
       )
       useStore.setState({ count: 0 }, false, 'reset')
+      useStore.persist.hasHydrated()
       return <></>
     }
     TestComponent
@@ -566,6 +593,7 @@ describe('more complex state spec with subscribeWithSelector', () => {
         (state) => state.foo,
         (foo) => console.log(foo)
       )
+      useStore.persist.hasHydrated()
       return <></>
     }
     TestComponent
diff --git a/tests/persistAsync.test.tsx b/tests/persistAsync.test.tsx
index dba5c52..592dfe7 100644
--- a/tests/persistAsync.test.tsx
+++ b/tests/persistAsync.test.tsx
@@ -451,7 +451,7 @@ describe('persist middleware with async configuration', () => {
       })
     )
     expect(useStore.persist.hasHydrated()).toBe(false)
-    await new Promise((resolve) => useStore.persist.onHydrate(resolve))
+    await new Promise((resolve) => useStore.persist.onFinishHydration(resolve))
     expect(useStore.persist.hasHydrated()).toBe(true)
 
     await useStore.persist.rehydrate()
diff --git a/tests/persistSync.test.tsx b/tests/persistSync.test.tsx
index beb2248..6ef1bb4 100644
--- a/tests/persistSync.test.tsx
+++ b/tests/persistSync.test.tsx
@@ -479,11 +479,14 @@ describe('persist middleware with sync configuration', () => {
     expect(useStore.persist.hasHydrated()).toBe(true)
   })
 
-  it('can wait for hydration through the api', async () => {
-    const storageValue = '{"state":{"count":1},"version":0}'
+  it('can wait for rehydration through the api', async () => {
+    const storageValue1 = '{"state":{"count":1},"version":0}'
+    const storageValue2 = '{"state":{"count":2},"version":0}'
 
     const onHydrateSpy1 = jest.fn()
     const onHydrateSpy2 = jest.fn()
+    const onFinishHydrationSpy1 = jest.fn()
+    const onFinishHydrationSpy2 = jest.fn()
 
     const storage = {
       getItem: () => '',
@@ -498,12 +501,29 @@ describe('persist middleware with sync configuration', () => {
       })
     )
 
-    useStore.persist.onHydrate(onHydrateSpy1)
+    const hydrateUnsub1 = useStore.persist.onHydrate(onHydrateSpy1)
     useStore.persist.onHydrate(onHydrateSpy2)
 
-    storage.getItem = () => storageValue
+    const finishHydrationUnsub1 = useStore.persist.onFinishHydration(
+      onFinishHydrationSpy1
+    )
+    useStore.persist.onFinishHydration(onFinishHydrationSpy2)
+
+    storage.getItem = () => storageValue1
+    await useStore.persist.rehydrate()
+    expect(onHydrateSpy1).toBeCalledWith({ count: 0 })
+    expect(onHydrateSpy2).toBeCalledWith({ count: 0 })
+    expect(onFinishHydrationSpy1).toBeCalledWith({ count: 1 })
+    expect(onFinishHydrationSpy2).toBeCalledWith({ count: 1 })
+
+    hydrateUnsub1()
+    finishHydrationUnsub1()
+
+    storage.getItem = () => storageValue2
     await useStore.persist.rehydrate()
-    expect(onHydrateSpy1).toBeCalledWith({ count: 1 })
+    expect(onHydrateSpy1).not.toBeCalledTimes(2)
     expect(onHydrateSpy2).toBeCalledWith({ count: 1 })
+    expect(onFinishHydrationSpy1).not.toBeCalledTimes(2)
+    expect(onFinishHydrationSpy2).toBeCalledWith({ count: 2 })
   })
 })