如何让Zustand自定义Hook支持类似useState的泛型类型传入？

需求背景 #

现有一个基于Zustand构建的自定义Hook，用于存储多种类型的数据，希望能像useState一样在调用时指定具体类型，确保TypeScript会对传入setActionData的非目标类型数据抛出错误，同时actionData的类型与指定类型匹配。

现有代码 #

import { Profile } from "@/types";
import { create } from "zustand";

type ActionData = Product | Category | Profile | null;

type ActionDataStore = {
    actionData: ActionData;
    setActionData: (actionData: ActionData) => void;
};

export const useActionData = create<ActionDataStore>((set) => ({
    actionData: null,
    setActionData: (actionData) =>
        set((state) => {
            return { actionData };
        }),
}));

改造方案 #

通过包装泛型函数的方式，在保留Zustand全局store特性的同时，实现调用时的类型约束：

import { Profile, Product, Category } from "@/types";
import { create } from "zustand";

// 定义所有支持的存储类型联合
type ActionData = Product | Category | Profile | null;

// 基础store类型定义
type ActionDataStore = {
  actionData: ActionData;
  setActionData: (actionData: ActionData) => void;
};

// 创建全局Zustand store
const useActionDataStore = create<ActionDataStore>((set) => ({
  actionData: null,
  setActionData: (actionData) => set({ actionData }),
}));

// 包装成泛型Hook对外暴露
export function useActionData<T extends ActionData>() {
  // 将返回的actionData断言为指定泛型类型
  const actionData = useActionDataStore((state) => state.actionData) as T;
  // 约束setActionData仅接受指定泛型类型的参数
  const setActionData = useActionDataStore((state) => (data: T) => state.setActionData(data));

  return { actionData, setActionData };
}

使用示例 #

在组件中调用时直接传入目标类型即可：

// 指定存储类型为Profile
const { actionData, setActionData } = useActionData<Profile>();

// 合法操作：传入Profile类型数据
setActionData({ id: 1, name: "Alice", email: "alice@example.com" });

// 非法操作：传入非Profile类型数据，TypeScript会抛出类型错误
setActionData({ id: 2, title: "Test Product", price: 99 });

此时actionData的类型为Profile | null，setActionData仅接受Profile | null类型的参数，完全符合需求。

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