make type NamedSet compatible with type SetState

package.json

@@ -142,6 +142,7 @@
     "eslint-plugin-react": "^7.24.0",
     "eslint-plugin-react-hooks": "^4.2.0",
     "husky": "^6.0.0",
+    "immer": "^9.0.3",
     "jest": "^27.0.4",
     "json": "^11.0.0",
     "lint-staged": "^11.0.0",

src/middleware.ts

@@ -31,8 +31,13 @@ export const redux =
   }
 
 export type NamedSet<T extends State> = {
-  <K extends keyof T>(
-    partial: PartialState<T, K>,
+  <
+    K1 extends keyof T,
+    K2 extends keyof T = K1,
+    K3 extends keyof T = K2,
+    K4 extends keyof T = K3
+  >(
+    partial: PartialState<T, K1, K2, K3, K4>,
     replace?: boolean,
     name?: string
   ): void

tests/middlewareTypes.test.tsx

@@ -0,0 +1,234 @@
+import { produce } from 'immer'
+import type { Draft } from 'immer'
+import create, { UseStore } from '../src'
+import { State, StateCreator } from '../src/vanilla'
+import { devtools, NamedSet, persist } from '../src/middleware'
+
+type TImmerConfigFn<T extends State> = (fn: (draft: Draft<T>) => void) => void
+type TImmerConfig<T extends State> = StateCreator<T, TImmerConfigFn<T>>
+
+interface ISelectors<T> {
+  use: {
+    [key in keyof T]: () => T[key]
+  }
+}
+
+const immer = <T extends State>(
+  config: TImmerConfig<T>
+): StateCreator<T, NamedSet<T>> => {
+  return (set, get, api) => {
+    return config((fn) => set(produce<T>(fn)), get, api)
+  }
+}
+
+const createSelectorHooks = <T extends State>(store: UseStore<T>) => {
+  const storeAsSelectors = store as unknown as ISelectors<T>
+  storeAsSelectors.use = {} as ISelectors<T>['use']
+
+  Object.keys(store.getState()).forEach((key) => {
+    const storeKey = key as keyof T
+    const selector = (state: T) => state[storeKey]
+
+    storeAsSelectors.use[storeKey] = () => store(selector)
+  })
+
+  return store as UseStore<T> & ISelectors<T>
+}
+
+interface ITestStateProps {
+  testKey: string
+  setTestKey: (testKey: string) => void
+}
+
+it('should have correct type when creating store with devtool', () => {
+  const createStoreWithDevtool = <T extends State>(
+    createState: StateCreator<T>,
+    prefix = 'prefix'
+  ): UseStore<T> & ISelectors<T> => {
+    return createSelectorHooks(create(devtools(createState, prefix)))
+  }
+
+  const testDevtoolStore = createStoreWithDevtool<ITestStateProps>(
+    (set) => ({
+      testKey: 'test',
+      setTestKey: (testKey: string) => {
+        set((state) => {
+          state.testKey = testKey
+        })
+      },
+    }),
+    'test'
+  )
+
+  const TestComponent = (): JSX.Element => {
+    testDevtoolStore.use.testKey()
+    testDevtoolStore.use.setTestKey()
+
+    return <></>
+  }
+})
+
+it('should have correct type when creating store with devtool and immer', () => {
+  const createStoreWithImmer = <T extends State>(
+    createState: TImmerConfig<T>,
+    prefix = 'prefix'
+  ): UseStore<T> & ISelectors<T> => {
+    return createSelectorHooks(create(devtools(immer(createState), prefix)))
+  }
+
+  const testImmerStore = createStoreWithImmer<ITestStateProps>(
+    (set) => ({
+      testKey: 'test',
+      setTestKey: (testKey: string) => {
+        set((state) => {
+          state.testKey = testKey
+        })
+      },
+    }),
+    'test'
+  )
+
+  const TestComponent = (): JSX.Element => {
+    testImmerStore.use.testKey()
+    testImmerStore.use.setTestKey()
+
+    return <></>
+  }
+})
+
+it('should have correct type when creating store with devtool and persist', () => {
+  const createStoreWithPersist = <T extends State>(
+    createState: StateCreator<T>,
+    prefix = 'prefix',
+    persistName = 'persist'
+  ): UseStore<T> & ISelectors<T> => {
+    return createSelectorHooks(
+      create(devtools(persist(createState, { name: persistName }), prefix))
+    )
+  }
+
+  const testPersistStore = createStoreWithPersist<ITestStateProps>(
+    (set) => ({
+      testKey: 'test',
+      setTestKey: (testKey: string) => {
+        set((state) => {
+          state.testKey = testKey
+        })
+      },
+    }),
+    'test',
+    'persist'
+  )
+
+  const TestComponent = (): JSX.Element => {
+    testPersistStore.use.testKey()
+    testPersistStore.use.setTestKey()
+
+    return <></>
+  }
+})
+
+it('should have correct type when creating store without middleware', () => {
+  const testStore = create<ITestStateProps>((set) => ({
+    testKey: 'test',
+    setTestKey: (testKey: string) => {
+      set((state) => {
+        state.testKey = testKey
+      })
+    },
+  }))
+
+  const TestComponent = (): JSX.Element => {
+    testStore((state) => state.testKey)
+    testStore((state) => state.setTestKey)
+
+    return <></>
+  }
+})
+
+it('should have correct type when creating store with persist', () => {
+  const createStoreWithPersist = <T extends State>(
+    createState: StateCreator<T>,
+    persistName = 'persist'
+  ): UseStore<T> & ISelectors<T> => {
+    return createSelectorHooks(
+      create(persist(createState, { name: persistName }))
+    )
+  }
+
+  const testPersistStore = createStoreWithPersist<ITestStateProps>(
+    (set) => ({
+      testKey: 'test',
+      setTestKey: (testKey: string) => {
+        set((state) => {
+          state.testKey = testKey
+        })
+      },
+    }),
+    'persist'
+  )
+
+  const TestComponent = (): JSX.Element => {
+    testPersistStore.use.testKey()
+    testPersistStore.use.setTestKey()
+
+    return <></>
+  }
+})
+
+it('should have correct type when creating store with immer', () => {
+  const createStoreWithImmer = <T extends State>(
+    createState: TImmerConfig<T>
+  ): UseStore<T> & ISelectors<T> => {
+    return createSelectorHooks(create(immer(createState)))
+  }
+
+  const testImmerStore = createStoreWithImmer<ITestStateProps>((set) => ({
+    testKey: 'test',
+    setTestKey: (testKey: string) => {
+      set((state) => {
+        state.testKey = testKey
+      })
+    },
+  }))
+
+  const TestComponent = (): JSX.Element => {
+    testImmerStore.use.testKey()
+    testImmerStore.use.setTestKey()
+
+    return <></>
+  }
+})
+
+it('should have correct type when creating store with devtool, persist and immer', () => {
+  const createStoreWithPersistAndImmer = <T extends State>(
+    createState: TImmerConfig<T>,
+    prefix = 'prefix',
+    persistName = 'persist'
+  ): UseStore<T> & ISelectors<T> => {
+    return createSelectorHooks(
+      create(
+        devtools(persist(immer(createState), { name: persistName }), prefix)
+      )
+    )
+  }
+
+  const testPersistImmerStore = createStoreWithPersistAndImmer<ITestStateProps>(
+    (set) => ({
+      testKey: 'test',
+      setTestKey: (testKey: string) => {
+        set((state) => {
+          state.testKey = testKey
+        })
+      },
+    }),
+    'test'
+  )
+
+  const TestComponent = (): JSX.Element => {
+    testPersistImmerStore.use.testKey()
+    testPersistImmerStore.use.setTestKey()
+
+    return <></>
+  }
+})

yarn.lock

@@ -3755,6 +3755,11 @@ ignore@^5.1.4:
   resolved "https://registry.yarnpkg.com/ignore/-/ignore-5.1.8.tgz#f150a8b50a34289b33e22f5889abd4d8016f0e57"
   integrity sha512-BMpfD7PpiETpBl/A6S498BaIJ6Y/ABT93ETbby2fP00v4EbvPBXWEoaR1UBPKs3iR53pJY7EtZk5KACI57i1Uw==
 
+immer@^9.0.3:
+  version "9.0.3"
+  resolved "https://registry.yarnpkg.com/immer/-/immer-9.0.3.tgz#146e2ba8b84d4b1b15378143c2345559915097f4"
+  integrity sha512-mONgeNSMuyjIe0lkQPa9YhdmTv8P19IeHV0biYhcXhbd5dhdB9HSK93zBpyKjp6wersSUgT5QyU0skmejUVP2A==
+
 import-fresh@^3.0.0, import-fresh@^3.2.1:
   version "3.3.0"
   resolved "https://registry.yarnpkg.com/import-fresh/-/import-fresh-3.3.0.tgz#37162c25fcb9ebaa2e6e53d5b4d88ce17d9e0c2b"