使用Immer中间件时为Zustand的set方法定义类型并保留自动补全

核心问题分析 #

启用Immer后，Zustand的set方法类型发生本质变化：回调参数从原状态类型变为Immer的WritableDraft<Store & ActionsObject>，且回调无需返回新状态（直接修改draft即可）。手动硬编码类型会丢失属性自动补全，复用非Immer版Setter类型则会触发类型不匹配错误。

解决方案 #

方案1：泛型函数自动推导类型（推荐）

让TypeScript自动从Zustand的Immer中间件中推导set和get的类型，避免手动硬编码：

import { create } from 'zustand';
import { immer } from 'zustand/middleware/immer';
import { WritableDraft } from 'immer';

// 定义状态与动作类型
type Store = {
  whateverProperty: string;
  // 其他状态属性
};

type Actions = {
  func: () => void;
  // 其他动作方法
};

// 用泛型约束set参数，确保回调接收WritableDraft类型的状态
function beginDraggingVertex<T extends Store & Actions>(
  set: (
    updater: ((state: WritableDraft<T>) => void) | Partial<T> | T,
    shouldReplace?: boolean
  ) => void,
  get: () => T,
  vertexID: string
) {
  return () => {
    set((store) => {
      // 此处自动补全store的所有属性
      store.whateverProperty = 'hi';
    });
  };
}

// 创建Store时传入泛型，自动推导内部set/get类型
export const useStore = create<Store & Actions>()(
  immer((set, get) => ({
    whateverProperty: 'initial',
    func: beginDraggingVertex(set, get, 'vertex-1')
  }))
);

方案2：手动定义正确的Immer版Setter类型

若需手动定义Setter类型，需引入Immer的WritableDraft并关联你的状态类型：

import { WritableDraft } from 'immer';

// 基于Store和Actions类型，定义Immer专属Setter
type ImmerSetter = (
  updater: 
    | (Store & Actions)
    | Partial<Store & Actions>
    | ((state: WritableDraft<Store & Actions>) => void),
  shouldReplace?: boolean
) => void;

type ImmerGetter = () => Store & Actions;

// 给外部函数指定正确类型
function beginDraggingVertex(set: ImmerSetter, get: ImmerGetter, vertexID: string) {
  return () => {
    set((store) => {
      // 自动补全恢复正常
      store.whateverProperty = 'hi';
    });
  };
}

方案3：结合Zustand内置类型工具

利用Zustand的StateSetter和Immer的Draft类型，快速构建正确的Setter类型：

import { StateSetter } from 'zustand';
import { Draft } from 'immer';

// 生成对应状态类型的Immer版Setter
type ImmerSetter = StateSetter<Draft<Store & Actions>>;

关键注意点 #

• 必须正确导入Immer的WritableDraft或Draft类型，否则TypeScript无法识别draft状态的属性
• 禁止混用非Immer版和Immer版的Setter类型，两者的回调逻辑与参数类型完全不同
• 优先依赖TypeScript自动推导类型，减少手动硬编码带来的类型不一致问题

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