zustand_immer_typescript.ts

import produce from 'immer';
import create, { GetState, SetState, State, StoreApi, UseStore } from 'zustand';

type StateCreator<T extends State, CustomSetState = SetState<T>, U extends State = T> = (
  set: CustomSetState,
  get: GetState<T>,
  api: StoreApi<T>
) => U;

const immer = <T extends State, U extends State>(
  config: StateCreator<T, (fn: (draft: T) => void) => void, U>
): StateCreator<T, SetState<T>, U> => (set, get, api) =>
  config((fn) => set(produce(fn) as (state: T) => T), get, api);

const combine = <PrimaryState extends State, SecondaryState extends State>(
  initialState: PrimaryState,
  create: (
    set: SetState<PrimaryState>,
    get: GetState<PrimaryState>,
    api: StoreApi<PrimaryState>
  ) => SecondaryState
): StateCreator<PrimaryState & SecondaryState> => (set, get, api) =>
  Object.assign(
    {},
    initialState,
    create(
      set as SetState<PrimaryState>,
      get as GetState<PrimaryState>,
      api as StoreApi<PrimaryState>
    )
  );

const combineAndImmer = <PrimaryState extends State, SecondaryState extends State>(
  initialState: PrimaryState,
  config: StateCreator<PrimaryState, (fn: (draft: PrimaryState) => void) => void, SecondaryState>
): StateCreator<PrimaryState & SecondaryState> => {
  return combine(initialState, immer(config));
};

/**
 * This is a custom store creator function.
 * It has a different api from the normal create.
 * But it provides typescript completion.
 *
 * @param initialState - The initialState of the store.
 * @param stateModifier - The set function given by zustand.
 *
 * Example:
 *
 * ```js
 * const initialState = { data: { count: 0, other: true } };
 *
 * const useStore = createStore(initialState, (set) => ({
 *   incBy: (by: number) => set((state) => (state.data.count += by)),
 *   invertOther: () => set((state) => (state.data.other = !state.data.other))
 * }));
 *
 * const date = useStore((state) => state.data);
 * ```
 */
const createStore = <PrimaryState extends State, SecondaryState extends State>(
  initialState: PrimaryState,
  config: StateCreator<PrimaryState, (fn: (draft: PrimaryState) => void) => void, SecondaryState>
): UseStore<PrimaryState & SecondaryState> => {
  return create(combineAndImmer(initialState, config));
};

export default createStore;