feat: allow return null from useMergeState updater function

react-restart · Mar 28, 2019 · c14a4e5 · c14a4e5
1 parent 5b749b6
commit c14a4e5
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/src/useMergeState.ts b/src/useMergeState.ts
@@ -1,20 +1,23 @@
 import { useState } from 'react'
 
-type Updater<TState> = (state: TState) => Partial<TState>
+type Updater<TState> = (state: TState) => Partial<TState> | null
 
 export type MergeStateSetter<TState> = (
   update: Updater<TState> | Partial<TState> | null
 ) => void
 
-export default function useMergeState<TState>(
+export default function useMergeState<TState extends {}>(
   initialState: TState
 ): [TState, MergeStateSetter<TState>] {
   const [state, setState] = useState<TState>(initialState)
 
   const updater = (update: Updater<TState> | Partial<TState> | null) => {
     if (update === null) return
     if (typeof update === 'function')
-      setState(state => ({ ...state, ...update(state) }))
+      setState(state => {
+        const nextState = update(state)
+        return nextState == null ? state : { ...state, ...nextState }
+      })
 
     setState(state => ({ ...state, ...update }))
   }