(types): improve types for setState (pmndrs#320)

* (types): improve types for setState

instead of using Partial, which would allow setting every field to undefined, we infer the types of the values depending on the keys passed into the set function

* (types): improve types for setState

add code comment to clarify why we are using Pick<T, K> | T

* (types): improve types for setState

apply the same changes we did for SetState also to namedSet for the devtools middleware

* (types): improve types for setState

tests for the partial logic

* empty commit

Co-authored-by: daishi <[email protected]>

src/middleware.ts

@@ -28,11 +28,13 @@ export const redux = <S extends State, A extends { type: unknown }>(
   return { dispatch: api.dispatch, ...initial }
 }
 
-type NamedSet<S extends State> = (
-  partial: PartialState<S>,
-  replace?: boolean,
-  name?: string
-) => void
+type NamedSet<T extends State> = {
+  <K extends keyof T>(
+    partial: PartialState<T, K>,
+    replace?: boolean,
+    name?: string
+  ): void
+}
 
 export const devtools = <S extends State>(
   fn: (set: NamedSet<S>, get: GetState<S>, api: StoreApi<S>) => S,
@@ -68,7 +70,10 @@ export const devtools = <S extends State>(
   const initialState = fn(namedSet, get, api)
   if (!api.devtools) {
     const savedSetState = api.setState
-    api.setState = (state: PartialState<S>, replace?: boolean) => {
+    api.setState = <K extends keyof S>(
+      state: PartialState<S, K>,
+      replace?: boolean
+    ) => {
       savedSetState(state, replace)
       api.devtools.send(api.devtools.prefix + 'setState', api.getState())
     }

src/vanilla.ts

@@ -1,7 +1,9 @@
 export type State = Record<string | number | symbol, unknown>
-export type PartialState<T extends State> =
-  | Partial<T>
-  | ((state: T) => Partial<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<T extends State, K extends keyof T = keyof T> =
+  | (Pick<T, K> | T)
+  | ((state: T) => Pick<T, K> | T)
 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
@@ -14,10 +16,10 @@ export interface Subscribe<T extends State> {
     equalityFn?: EqualityChecker<StateSlice>
   ): () => void
 }
-export type SetState<T extends State> = (
-  partial: PartialState<T>,
-  replace?: boolean
-) => void
+
+export type SetState<T extends State> = {
+  <K extends keyof T>(partial: PartialState<T, K>, replace?: boolean): void
+}
 export type GetState<T extends State> = () => T
 export type Destroy = () => void
 export interface StoreApi<T extends State> {

tests/basic.test.tsx

@@ -46,8 +46,13 @@ it('creates a store hook and api object', () => {
   `)
 })
 
+type CounterState = {
+  count: number
+  inc: () => void
+}
+
 it('uses the store with no args', async () => {
-  const useStore = create<any>((set) => ({
+  const useStore = create<CounterState>((set) => ({
     count: 0,
     inc: () => set((state) => ({ count: state.count + 1 })),
   }))
@@ -66,7 +71,7 @@ it('uses the store with no args', async () => {
 it('uses the store with selectors', async () => {
   const useStore = create<any>((set) => ({
     count: 0,
-    inc: () => set((state) => ({ count: state.count + 1 })),
+    inc: () => set((state: any) => ({ count: state.count + 1 })),
   }))
 
   function Counter() {
@@ -115,7 +120,7 @@ it('uses the store with a selector and equality checker', async () => {
 it('only re-renders if selected state has changed', async () => {
   const useStore = create<any>((set) => ({
     count: 0,
-    inc: () => set((state) => ({ count: state.count + 1 })),
+    inc: () => set((state: any) => ({ count: state.count + 1 })),
   }))
   let counterRenderCount = 0
   let controlRenderCount = 0
@@ -167,7 +172,7 @@ it('re-renders with useLayoutEffect', async () => {
 it('can batch updates', async () => {
   const useStore = create<any>((set) => ({
     count: 0,
-    inc: () => set((state) => ({ count: state.count + 1 })),
+    inc: () => set((state: any) => ({ count: state.count + 1 })),
   }))
 
   function Counter() {
@@ -742,3 +747,35 @@ it('can use exposed types', () => {
     useStore
   )
 })
+
+type AssertEqual<Type, Expected> = Type extends Expected
+  ? Expected extends Type
+    ? true
+    : never
+  : never
+
+it('should have correct (partial) types for setState', () => {
+  type Count = { count: number }
+
+  const store = create<Count>((set) => ({
+    count: 0,
+    // @ts-expect-error we shouldn't be able to set count to undefined
+    a: () => set(() => ({ count: undefined })),
+    // @ts-expect-error we shouldn't be able to set count to undefined
+    b: () => set({ count: undefined }),
+    c: () => set({ count: 1 }),
+  }))
+
+  const setState: AssertEqual<typeof store.setState, SetState<Count>> = true
+  expect(setState).toEqual(true)
+
+  // ok, should not error
+  store.setState({ count: 1 })
+  store.setState({})
+  store.setState(() => {})
+
+  // @ts-expect-error type undefined is not assignable to type number
+  store.setState({ count: undefined })
+  // @ts-expect-error type undefined is not assignable to type number
+  store.setState((state) => ({ ...state, count: undefined }))
+})