make devtools replace-friendly and more correctly typed

pmndrs · Apr 12, 2022 · 58dc684 · 58dc684
1 parent f0d192a
commit 58dc684
Show file tree

Hide file tree

Showing 2 changed files with 107 additions and 5 deletions.
diff --git a/src/middleware/devtools.ts b/src/middleware/devtools.ts
@@ -1,4 +1,5 @@
 import type {} from '@redux-devtools/extension'
+import type { Draft } from 'immer'
 import {
   PartialState,
   SetState,
@@ -24,15 +25,62 @@ type Message = {
 
 type Write<T extends object, U extends object> = Omit<T, keyof U> & U
 type Cast<T, U> = T extends U ? T : U
+type TakeTwo<T> = T extends []
+  ? [undefined, undefined]
+  : T extends [unknown]
+  ? [...a0: T, a1: undefined]
+  : T extends [unknown?]
+  ? [...a0: T, a1: undefined]
+  : T extends [unknown, unknown]
+  ? T
+  : T extends [unknown, unknown?]
+  ? T
+  : T extends [unknown?, unknown?]
+  ? T
+  : T extends [infer A0, infer A1, ...unknown[]]
+  ? [A0, A1]
+  : T extends [infer A0, (infer A1)?, ...unknown[]]
+  ? [A0, A1?]
+  : T extends [(infer A0)?, (infer A1)?, ...unknown[]]
+  ? [A0?, A1?]
+  : never
 
 type WithDevtools<S> = Write<Cast<S, object>, StoreSetStateWithAction<S>>
 
 type StoreSetStateWithAction<S> = S extends {
-  setState: (...a: infer A) => infer R
+  getState: () => infer T
+  setState: (...a: infer A) => unknown
 }
-  ? {
-      setState: (...a: [...a: A, actionType?: string | { type: unknown }]) => R
-    }
+  ? A extends [Partial<T> | ((state: T) => Partial<T>), (boolean | undefined)?]
+    ? {
+        setState: <
+          Nt extends R extends true ? T : Partial<T>,
+          R extends boolean | undefined
+        >(
+          partial: Nt | ((state: T) => Nt),
+          replace?: R,
+          actionType?: string | { type: unknown }
+        ) => void
+      }
+    : A extends [
+        Partial<T> | ((state: Draft<T>) => void),
+        (boolean | undefined)?
+      ]
+    ? {
+        setState: <
+          Nt extends R extends true ? T : Partial<T>,
+          R extends boolean | undefined
+        >(
+          nextStateOrUpdater: Nt | ((state: Draft<T>) => void),
+          shouldReplace?: R,
+          actionType?: string | { type: unknown }
+        ) => void
+      }
+    : {
+        setState: (
+          ...a: [...a: TakeTwo<A>, actionType?: string | { type: unknown }]
+        ) => void
+      }
   : never
 
 interface DevtoolsOptions {

diff --git a/tests/types.test.tsx b/tests/types.test.tsx
@@ -12,6 +12,7 @@ import create, {
   Subscribe,
   UseBoundStore,
 } from 'zustand'
+import { devtools, immer } from 'zustand/middleware'
 
 it('can use exposed types', () => {
   interface ExampleState {
@@ -174,7 +175,7 @@ it('should allow for different partial keys to be returnable from setState', ()
 })
 
 it('setState with replace requires whole state', () => {
-  const store = create<{ count: number; something: string }>(() => ({
+  const store = create<{ count: number; something: string }>()(() => ({
     count: 0,
     something: 'foo',
   }))
@@ -184,4 +185,57 @@ it('setState with replace requires whole state', () => {
     { count: 1 },
     true
   )
+
+  const storeWithDevtools = create<{ count: number; something: string }>()(
+    devtools(() => ({
+      count: 0,
+      something: 'foo',
+    }))
+  )
+
+  storeWithDevtools.setState(
+    // @ts-expect-error missing `something`
+    { count: 1 },
+    true
+  )
+
+  storeWithDevtools.setState(
+    // @ts-expect-error missing `something`
+    { count: 1 },
+    true,
+    'FOO'
+  )
+
+  storeWithDevtools.setState({ count: 1 }, false, 'FOO')
+
+  storeWithDevtools.setState({ count: 1 }, false)
+
+  const storeWithDevtoolsAndImmer = create<{
+    count: number
+    something: string
+  }>()(
+    immer(
+      devtools(() => ({
+        count: 0,
+        something: 'foo',
+      }))
+    )
+  )
+
+  storeWithDevtoolsAndImmer.setState(
+    // @ts-expect-error missing `something`
+    { count: 1 },
+    true
+  )
+
+  storeWithDevtoolsAndImmer.setState(
+    // @ts-expect-error missing `something`
+    { count: 1 },
+    true,
+    'FOO'
+  )
+
+  storeWithDevtoolsAndImmer.setState({ count: 1 }, false, 'FOO')
+
+  storeWithDevtoolsAndImmer.setState({ count: 1 }, false)
 })