一种组织 zustand 自定义函数的方法

问题

我们在使用 zustand 的时候，总会遇到需要自定义函数：

// 定义
const useCounter = create((set) => ({
  count: 0,
  inc: () => {
    set((prev) => ({
      count: prev.count + 1,
    }))
  },
}))

// 使用
function App() {
  const { count, inc } = useCounter()

  return <button
    onClick={inc}
  >
    {count}
  </button>
}

这样有几个问题：

1. 函数污染了 store，我们在使用 useCounter 的时候，inc 函数总是在提示列表里面碍眼；
2. 当我们需要在 hook 里面使用函数时，平白增加了依赖

   function App() {
     const { count, inc } = useCounter()
   
     useEffect(() => {
       console.log(count)
       inc()
       // inc 是一个静态函数，按理来说是可以不用放在依赖里面的
     }, [count, inc])
   }
   

3. 没有类型限制，无法阻止别人无意间修改 inc 函数：

   // 在一些场合下，我们希望使用 replaceFlag，这时候就出问题了，把 inc 函数丢了
   // replace flag: https://docs.pmnd.rs/zustand/guides/immutable-state-and-merging#replace-flag
   const replaceFlag = true
   useCounter.setState({
     count: 3,
   }, replaceFlag)
   

这个怎么样？

如果能这样呢：

function App() {
  const { count } = useCounter()

  return <button
    // 将 inc 变成 useCounter 的静态函数
    onClick={useCounter.inc}
  >
    {count}
  </button>
}

直接上代码

/* eslint-disable @typescript-eslint/no-explicit-any */
import type { StoreApi, UseBoundStore } from 'zustand'

interface StaticFuncs {
  [key: string]: (...args: any[]) => any
}

type WithStatic<
  T extends UseBoundStore<StoreApi<any>>,
  S extends StaticFuncs
> = T & S

export function withStatic<T, S extends StaticFuncs>(
  useStore: UseBoundStore<StoreApi<T>>,
  staticFuncs: S
) {
  const protectedKeys = Object.keys(useStore)
  const result = useStore as WithStatic<UseBoundStore<StoreApi<T>>, S>
  Object.keys(staticFuncs).forEach((key) => {
    if (protectedKeys.includes(key)) {
      if (isDev) {
        throw new Error(`protected key: ${key}`)
      } else {
        console.error(`protected key: ${key}`)
      }
      return
    }
    result[key as 'setState'] = staticFuncs[key]
  })
  return result
}

Usage

const useRawCounter = create(() => ({
  count: 0,
}))

export const useCounter = withStatic(useRawCounter, {
  inc: () => {
    useRawCounter.setState((prev) => ({
      count: prev.count + 1,
    }))
  },
})

// 调用
function App() {
  const { count } = useCounter()

  return <button
    onClick={useCounter.inc}
  >
    {count}
  </button>
}