如何在Zustand切片模式下为切片添加根属性且不修改现有代码？

问题背景 #

我正在使用Zustand的切片模式，现有代码结构如下：

import { create, StateCreator } from 'zustand'

interface BearSlice {
  bears: number
  addBear: () => void
  eatFish: () => void
}

interface FishSlice {
  fishes: number
  addFish: () => void
}

interface SharedSlice {
  addBoth: () => void
  getBoth: () => void
}

const createBearSlice: StateCreator<BearSlice & FishSlice, [], [], BearSlice> = (set) => ({
  bears: 0,
  addBear: () => set((state) => ({ bears: state.bears + 1 })),
  eatFish: () => set((state) => ({ fishes: state.fishes - 1 })),
})

const createFishSlice: StateCreator<BearSlice & FishSlice, [], [], FishSlice> = (set) => ({
  fishes: 0,
  addFish: () => set((state) => ({ fishes: state.fishes + 1 })),
})

const createSharedSlice: StateCreator<BearSlice & FishSlice, [], [], SharedSlice> = (set, get) => ({
  addBoth: () => {
    // 可以复用已有方法
    get().addBear()
    get().addFish()
    // 也可以直接修改状态
    // set((state) => ({ bears: state.bears + 1, fishes: state.fishes + 1 })
  },
  getBoth: () => get().bears + get().fishes,
})

const useBoundStore = create<BearSlice & FishSlice & SharedSlice>()((...a) => ({
  ...createBearSlice(...a),
  ...createFishSlice(...a),
  ...createSharedSlice(...a),
}))

我希望给每个切片添加根属性包裹，让Store结构变成这样：

const useBoundStore = create<{ bearSlice: BearSlice; fishSlice: FishSlice, sharedSlice: SharedSlice }>()((...a) => ({
  bearSlice: createBearSlice(...a),
  fishSlice: createFishSlice(...a),
  sharedSlice: createSharedSlice(...a),
}))

但手动修改所有切片内部的set语句工作量极大，比如要把：

set((state) => ({ bears: state.bears + 1 }))

改成：

set((state) => ({ bearSlice: { ...state.bearSlice, bears: state.bearSlice.bears + 1 }} ))

以下是几种可行的解决方案：

方案1：封装切片包装函数，自动映射根属性 #

创建通用包装函数，帮切片自动处理根属性的set和get，原切片代码无需修改：

import { create, StateCreator } from 'zustand'

// 定义切片包装类型
type SliceWrapper<T extends object, K extends string> = StateCreator<
  Record<K, T>,
  [],
  [],
  T
>

// 生成带根属性映射的切片创建函数
const wrapSlice = <T extends object, K extends string>(
  sliceKey: K,
  createSlice: StateCreator<T, [], [], T>
): SliceWrapper<T, K> => (set, get, api) => {
  // 包装set，自动添加根属性前缀
  const wrappedSet: typeof set = (updater) => {
    if (typeof updater === 'function') {
      return set((rootState) => {
        const sliceState = rootState[sliceKey]
        const updatedSlice = updater(sliceState)
        return { [sliceKey]: { ...sliceState, ...updatedSlice } }
      })
    } else {
      return set({ [sliceKey]: { ...get()[sliceKey], ...updater } })
    }
  }

  // 包装get，直接返回当前切片状态
  const wrappedGet: typeof get = () => get()[sliceKey]

  return createSlice(wrappedSet, wrappedGet, api)
}

// 原切片定义保持不变
interface BearSlice {
  bears: number
  addBear: () => void
  eatFish: () => void
}

interface FishSlice {
  fishes: number
  addFish: () => void
}

interface SharedSlice {
  addBoth: () => void
  getBoth: () => void
}

// 用包装函数处理原切片
const createBearSlice = wrapSlice('bearSlice', (set) => ({
  bears: 0,
  addBear: () => set((state) => ({ bears: state.bears + 1 })),
}))

const createFishSlice = wrapSlice('fishSlice', (set) => ({
  fishes: 0,
  addFish: () => set((state) => ({ fishes: state.fishes + 1 })),
}))

// 共享切片单独处理，需访问完整根状态
const createSharedSlice: StateCreator<
  { bearSlice: BearSlice; fishSlice: FishSlice },
  [],
  [],
  SharedSlice
> = (set, get) => ({
  addBoth: () => {
    get().bearSlice.addBear()
    get().fishSlice.addFish()
  },
  getBoth: () => get().bearSlice.bears + get().fishSlice.fishes,
})

// 创建最终Store
const useBoundStore = create<{
  bearSlice: BearSlice
  fishSlice: FishSlice
  sharedSlice: SharedSlice
}>()((...a) => ({
  bearSlice: createBearSlice(...a),
  fishSlice: createFishSlice(...a),
  sharedSlice: createSharedSlice(...a),
}))

注意：如果切片内有跨切片修改逻辑（比如原BearSlice的eatFish），此方案无法直接支持，需改用方案2。

方案2：自定义命名空间中间件，兼容跨切片操作 #

通过Zustand中间件+Proxy代理，让切片代码几乎无需改动即可支持根属性包裹：

import { create, StateCreator, Middleware } from 'zustand'

// 生成命名空间中间件
const withNamespace = <K extends string>(namespace: K): Middleware => (
  createState
) => (set, get, api) => {
  // 包装set，自动映射到对应根属性
  const namespaceSet = (updater: any) => {
    if (typeof updater === 'function') {
      return set((rootState) => {
        const sliceState = rootState[namespace]
        const update = updater(sliceState)
        return { [namespace]: { ...sliceState, ...update } }
      })
    } else {
      if (updater[namespace]) return set(updater)
      return set({ [namespace]: { ...get()[namespace], ...updater } })
    }
  }

  // 包装get，通过Proxy实现切片状态与根状态的无缝访问
  const namespaceGet = () => {
    const rootState = get()
    return new Proxy(rootState, {
      get(target, prop) {
        if (prop in target[namespace]) {
          return target[namespace][prop as keyof typeof target[namespace]]
        }
        return target[prop as keyof typeof target]
      },
    }) as typeof rootState & typeof rootState[K]
  }

  return createState(namespaceSet, namespaceGet, api)
}

// 原切片定义不变
interface BearSlice {
  bears: number
  addBear: () => void
  eatFish: () => void
}

interface FishSlice {
  fishes: number
  addFish: () => void
}

interface SharedSlice {
  addBoth: () => void
  getBoth: () => void
}

// 给切片添加命名空间支持
const createBearSlice: StateCreator<
  { bearSlice: BearSlice; fishSlice: FishSlice },
  [],
  [],
  BearSlice
> = withNamespace('bearSlice')((set) => ({
  bears: 0,
  addBear: () => set((state) => ({ bears: state.bears + 1 })),
  eatFish: () => set((state) => ({ fishes: state.fishes - 1 })), // 仍可直接跨切片修改
}))

const createFishSlice: StateCreator<
  { bearSlice: BearSlice; fishSlice: FishSlice },
  [],
  [],
  FishSlice
> = withNamespace('fishSlice')((set) => ({
  fishes: 0,
  addFish: () => set((state) => ({ fishes: state.fishes + 1 })),
}))

// 共享切片逻辑不变
const createSharedSlice: StateCreator<
  { bearSlice: BearSlice; fishSlice: FishSlice },
  [],
  [],
  SharedSlice
> = (set, get) => ({
  addBoth: () => {
    get().bearSlice.addBear()
    get().fishSlice.addFish()
  },
  getBoth: () => get().bearSlice.bears + get().fishSlice.fishes,
})

// 创建最终Store
const useBoundStore = create<{
  bearSlice: BearSlice
  fishSlice: FishSlice
  sharedSlice: SharedSlice
}>()((...a) => ({
  bearSlice: createBearSlice(...a),
  fishSlice: createFishSlice(...a),
  sharedSlice: createSharedSlice(...a),
}))

此方案通过代理实现了切片内部对状态的无缝访问，跨切片修改逻辑也能保留，几乎不需要改动原有切片代码。

方案3：IDE批量替换+类型推导（适合小体量代码库） #

如果切片数量不多，可以用IDE的正则批量替换快速完成修改：

1. 更新Store的类型定义为带根属性的结构；
2. 用正则替换将set\((state) => \({替换为set((state) => ({ bearSlice: { ...state.bearSlice,（对应BearSlice），同理处理FishSlice；
3. 对跨切片修改的语句单独调整，比如将eatFish的修改改为set((state) => ({ fishSlice: { ...state.fishSlice, fishes: state.fishSlice.fishes - 1 } }))。

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