Improve type definition of create to accept CustomSetState

src/index.ts

@@ -30,8 +30,11 @@ export interface UseStore<T extends State> {
   destroy: Destroy
 }
 
-export default function create<TState extends State>(
-  createState: StateCreator<TState> | StoreApi<TState>
+export default function create<
+  TState extends State,
+  CustomSetState extends SetState<TState> = SetState<TState>
+>(
+  createState: StateCreator<TState, CustomSetState> | StoreApi<TState>
 ): UseStore<TState> {
   const api: StoreApi<TState> =
     typeof createState === 'function' ? createImpl(createState) : createState

src/vanilla.ts

@@ -1,4 +1,18 @@
 export type State = object
+type PartialStatePlain<
+  T extends State,
+  K1 extends keyof T = keyof T,
+  K2 extends keyof T = K1,
+  K3 extends keyof T = K2,
+  K4 extends keyof T = K3
+> = Pick<T, K1> | Pick<T, K2> | Pick<T, K3> | Pick<T, K4> | T
+type PartialStateCallable<
+  T extends State,
+  K1 extends keyof T = keyof T,
+  K2 extends keyof T = K1,
+  K3 extends keyof T = K2,
+  K4 extends keyof T = K3
+> = (state: T) => Pick<T, K1> | Pick<T, K2> | Pick<T, K3> | Pick<T, K4> | T
 // types inspired by setState from React, see:
 // https://github.com/DefinitelyTyped/DefinitelyTyped/blob/6c49e45842358ba59a508e13130791989911430d/types/react/v16/index.d.ts#L489-L495
 export type PartialState<
@@ -8,8 +22,9 @@ export type PartialState<
   K3 extends keyof T = K2,
   K4 extends keyof T = K3
 > =
-  | (Pick<T, K1> | Pick<T, K2> | Pick<T, K3> | Pick<T, K4> | T)
-  | ((state: T) => Pick<T, K1> | Pick<T, K2> | Pick<T, K3> | Pick<T, K4> | T)
+  | PartialStatePlain<T, K1, K2, K3, K4>
+  | PartialStateCallable<T, K1, K2, K3, K4>
+
 export type StateSelector<T extends State, U> = (state: T) => U
 export type EqualityChecker<T> = (state: T, newState: T) => boolean
 export type StateListener<T> = (state: T, previousState: T) => void
@@ -48,19 +63,15 @@ export type StateCreator<T extends State, CustomSetState = SetState<T>> = (
   api: StoreApi<T>
 ) => T
 
-export default function create<TState extends State>(
-  createState: StateCreator<TState>
-): StoreApi<TState> {
+export default function create<
+  TState extends State,
+  CustomSetState extends SetState<TState> = SetState<TState>
+>(createState: StateCreator<TState, CustomSetState>): StoreApi<TState> {
   let state: TState
   const listeners: Set<StateListener<TState>> = new Set()
 
   const setState: SetState<TState> = (partial, replace) => {
-    // TODO: Remove type assertion once https://github.com/microsoft/TypeScript/issues/37663 is resolved
-    // https://github.com/microsoft/TypeScript/issues/37663#issuecomment-759728342
-    const nextState =
-      typeof partial === 'function'
-        ? (partial as (state: TState) => TState)(state)
-        : partial
+    const nextState = typeof partial === 'function' ? partial(state) : partial
     if (nextState !== state) {
       const previousState = state
       state = replace

tests/basic.test.tsx

@@ -395,9 +395,14 @@ it('can set the store', () => {
 })
 
 it('can set the store without merging', () => {
-  const { setState, getState } = create((_set) => ({
-    a: 1,
-  }))
+  type State = { a?: number; b?: number }
+
+  const { setState, getState } = create(
+    (_set) =>
+      ({
+        a: 1,
+      } as State)
+  )
 
   // Should override the state instead of merging.
   setState({ b: 2 }, true)

tests/types.test.tsx

@@ -173,3 +173,39 @@ it('should allow for different partial keys to be returnable from setState', ()
     return { something: true }
   })
 })
+
+it('can accept a store has CustomSetState', () => {
+  function dummyMiddleware<T extends State, Settable extends keyof T>(
+    creator: StateCreator<T, SetState<Pick<T, Settable>>>
+  ) {
+    return creator as unknown as StateCreator<T>
+  }
+
+  type ExampleState = {
+    a: string
+    b: string
+    readA(): string
+    readB(): string
+    send(): void
+  }
+
+  const store = create(
+    dummyMiddleware<ExampleState, 'a' | 'b'>((set, get) => ({
+      a: 'aaa',
+      b: 'bbbb',
+      readA() {
+        return get().a
+      },
+      readB() {
+        return get().b
+      },
+      send() {
+        // @ts-expect-error  Argument of type '{ readA: () => string; }' is not assignable to parameter of type 'PartialState<Pick<ExampleState, "a" | "b">, "a" | "b", "a" | "b", "a" | "b", "a" | "b">'.
+        set({ readA: () => 'ttt' })
+      },
+    }))
+  )
+
+  store.setState({ a: 'next a' })
+  store.setState({ readA: () => 'next a' })
+})