为基于Car ID创建的Zustand Stores实现类型化重置函数

问题背景 #

我通过Car ID创建了多个Zustand Store，需要一个能一键重置数十个Store的函数。找到了一个无类型的实现方案，但无法正确完成类型定义。

原无类型实现 #

const resetSet = new Set();
export const resettable = (create: any) => (initialState: any) => {
  const useStore = create(() => initialState);
  const reset = () => useStore.setState(initialState, true);
  resetSet.add(reset);
  return useStore;
};
export const resetAll = () => resetSet.forEach((reset: any) => reset());

示例用法 #

export type CarsState = {
  carInfo: CarInfoType | null;
  setCarInfo: (carInfo: CarInfoType | null) => void;
  // ...其他状态和方法
};

export const useCarState = (carID: string) =>
  resettable(
    create<CarState>()(
      persist(immer(handler), {
        name: `carState-${carID}`,
        storage: createJSONStorage(() => sessionStorage),
      }),
    ),
  );

类型化尝试的问题 #

我尝试添加类型定义，但错误地将useStore当成StateCreator调用（实际useStore是Zustand的钩子实例，拥有setState方法），导致类型不匹配。尝试代码如下：

import { StateCreator } from 'zustand';

type CreateFunction = <StateType>(initializer: () => StateType) => StateCreator<StateType>;

type ResettableFunction = <T>(create: CreateFunction) => (initialState: T) => StateCreator<T>;

const resetSet = new Set<() => void>();

export const resettable: ResettableFunction = <T>(create: CreateFunction) => (initialState: T) => {
  const useStore = create(() => initialState);
  const reset = () => useStore(initialState, true); // 错误：useStore不是函数，应该调用useStore.setState
  resetSet.add(reset);
  return useStore;
};

正确的类型化实现 #

核心是正确识别Zustand的UseBoundStore类型，它包含setState方法。以下是完整的类型安全实现：

import { UseBoundStore, StateCreator } from 'zustand';

// 存储所有重置函数的集合
const resetSet = new Set<() => void>();

/**
 * 创建可重置的Zustand Store高阶函数
 * @param create - Zustand的create函数（或经过中间件包装后的create结果）
 * @returns 接收初始状态的函数，返回带重置能力的Store
 */
export const resettable = <T>(
  create: (stateCreator: StateCreator<T>) => UseBoundStore<StateCreator<T>>
) => {
  return (initialState: T) => {
    // 创建Store实例
    const useStore = create(() => initialState);

    // 定义重置函数：替换整个状态为初始值
    const reset = () => {
      useStore.setState(initialState, true);
    };

    // 将重置函数加入集合
    resetSet.add(reset);

    return useStore;
  };
};

/**
 * 一键重置所有通过resettable创建的Store
 */
export const resetAll = () => {
  resetSet.forEach(reset => reset());
};

如果你的场景中，create参数是直接传入Zustand的create柯里化版本（比如create<CarState>()），可以进一步优化类型，让TS自动推导状态类型：

import { UseBoundStore, StateCreator, CreateState } from 'zustand';

const resetSet = new Set<() => void>();

export const resettable = <T>(
  create: CreateState<T>
) => {
  return (initialState: T) => {
    const useStore = create(() => initialState);
    const reset = () => useStore.setState(initialState, true);
    resetSet.add(reset);
    return useStore;
  };
};

export const resetAll = () => resetSet.forEach(reset => reset());

优化后的示例用法 #

export type CarsState = {
  carInfo: CarInfoType | null;
  setCarInfo: (carInfo: CarInfoType | null) => void;
};

// 定义状态处理器
const handler: StateCreator<CarsState> = (set) => ({
  carInfo: null,
  setCarInfo: (info) => set({ carInfo: info }),
});

export const useCarState = (carID: string) => {
  // 初始状态可以直接从handler推导，或显式传入
  const initialState: CarsState = { carInfo: null, setCarInfo: () => {} };
  
  return resettable(
    create<CarsState>()(
      persist(immer(handler), {
        name: `carState-${carID}`,
        storage: createJSONStorage(() => sessionStorage),
      })
    )(initialState); // 这里传入初始状态
};

内容的提问来源于stack exchange，提问作者Droidas