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