| 1 | import { create } from 'zustand'; |
| 2 | import { combine, persist, createJSONStorage } from 'zustand/middleware'; |
| 3 | import { deepClone } from './clone'; |
| 4 | import { indexedDBStorage } from './storage'; |
| 5 | |
| 6 | type Updater<T> = (updater: (value: T) => void) => void; |
| 7 | |
| 8 | type SecondParam<T> = T extends ( |
| 9 | _f: infer _F, |
| 10 | _s: infer S, |
| 11 | ...args: infer _U |
| 12 | ) => any |
| 13 | ? S |
| 14 | : never; |
| 15 | |
| 16 | type MakeUpdater<T> = { |
| 17 | lastUpdateTime: number; |
| 18 | _hasHydrated: boolean; |
| 19 | |
| 20 | markUpdate: () => void; |
| 21 | update: Updater<T>; |
| 22 | setHasHydrated: (state: boolean) => void; |
| 23 | }; |
| 24 | |
| 25 | type SetStoreState<T> = ( |
| 26 | partial: T | Partial<T> | ((state: T) => T | Partial<T>), |
| 27 | replace?: boolean | undefined, |
| 28 | ) => void; |
| 29 | |
| 30 | /** |
| 31 | * 创建一个数据持久化的 store |
| 32 | * 想要更改的时候持久化数据必须使用 set 方法 |
| 33 | * @param state |
| 34 | * @param methods |
| 35 | * @param persistOptions |
| 36 | */ |
| 37 | export function createPersistStore<T extends object, M>( |
| 38 | state: T, |
| 39 | methods: ( |
| 40 | set: SetStoreState<T & MakeUpdater<T>>, |
| 41 | get: () => T & MakeUpdater<T>, |
| 42 | ) => M, |
| 43 | persistOptions: SecondParam<typeof persist<T & M & MakeUpdater<T>>>, |
| 44 | ) { |
| 45 | persistOptions.storage = createJSONStorage(() => indexedDBStorage); |
| 46 | const oldOonRehydrateStorage = persistOptions?.onRehydrateStorage; |
| 47 | persistOptions.onRehydrateStorage = (state) => { |
| 48 | oldOonRehydrateStorage?.(state); |
| 49 | return () => state.setHasHydrated(true); |
| 50 | }; |
| 51 | |
| 52 | return create( |
| 53 | persist( |
| 54 | combine( |
| 55 | { |
| 56 | ...state, |
| 57 | lastUpdateTime: 0, |
| 58 | _hasHydrated: false, |
| 59 | }, |
| 60 | (set, get) => { |
| 61 | return { |
| 62 | ...methods(set as any, get as any), |
| 63 | |
| 64 | markUpdate() { |
| 65 | set({ lastUpdateTime: Date.now() } as Partial< |
| 66 | T & M & MakeUpdater<T> |
| 67 | >); |
| 68 | }, |
| 69 | update(updater) { |
| 70 | const state = deepClone(get()); |
| 71 | updater(state); |
| 72 | set({ |
| 73 | ...state, |
| 74 | lastUpdateTime: Date.now(), |
| 75 | }); |
| 76 | }, |
| 77 | setHasHydrated: (state: boolean) => { |
| 78 | set({ _hasHydrated: state } as Partial<T & M & MakeUpdater<T>>); |
| 79 | }, |
| 80 | } as M & MakeUpdater<T>; |
| 81 | }, |
| 82 | ), |
| 83 | persistOptions as any, |
| 84 | ), |
| 85 | ); |
| 86 | } |
| 87 |