feat: add useAsyncSetState (#12)

* feat: add useAsyncSetState

* clean up async state

Author: jquense

.babelrc

@@ -1,5 +1,5 @@
 {
-  "presets": [["@4c"], "@babel/typescript"],
+  "presets": ["@4c", "@babel/typescript"],
   "env": {
     "esm": {
       "presets": [
@@ -10,6 +10,9 @@
           }
         ]
       ]
+    },
+    "test": {
+      "presets": [["@4c", { "development": true }]]
     }
   }
 }

package.json

@@ -22,7 +22,7 @@
   },
   "jest": {
     "preset": "@4c",
-    "setupFiles": [
+    "setupFilesAfterEnv": [
       "./test/setup.js"
     ]
   },

src/useAnimationFrame.ts

@@ -1,7 +1,7 @@
 import { useRef } from 'react'
-import useWillUnmount from './useWillUnmount'
 import useMounted from './useMounted'
 import useStableMemo from './useStableMemo'
+import useWillUnmount from './useWillUnmount'
 
 export interface UseAnimationFrameReturn {
   cancel(): void
@@ -22,7 +22,7 @@ export interface UseAnimationFrameReturn {
  * Returns a controller object for requesting and cancelling an animation freame that is properly cleaned up
  * once the component unmounts. New requests cancel and replace existing ones.
  *
- * ```tsx
+ * ```ts
  * const [style, setStyle] = useState({});
  * const animationFrame = useAnimationFrame();
  *

src/useMergeState.ts

@@ -1,4 +1,4 @@
-import { useState } from 'react'
+import { useCallback, useState } from 'react'
 
 type Updater<TState> = (state: TState) => Partial<TState> | null
 
@@ -29,17 +29,20 @@ export default function useMergeState<TState extends {}>(
 ): [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 => {
-        const nextState = update(state)
-        return nextState == null ? state : { ...state, ...nextState }
-      })
-    } else {
-      setState(state => ({ ...state, ...update }))
-    }
-  }
+  const updater = useCallback(
+    (update: Updater<TState> | Partial<TState> | null) => {
+      if (update === null) return
+      if (typeof update === 'function') {
+        setState(state => {
+          const nextState = update(state)
+          return nextState == null ? state : { ...state, ...nextState }
+        })
+      } else {
+        setState(state => ({ ...state, ...update }))
+      }
+    },
+    [setState],
+  )
 
   return [state, updater]
 }

src/useStateAsync.ts

@@ -0,0 +1,71 @@
+import React, { useCallback, useEffect, useRef, useState } from 'react'
+
+type Updater<TState> = (state: TState) => TState
+
+export type AsyncSetState<TState> = (
+  stateUpdate: React.SetStateAction<TState>,
+) => Promise<TState>
+
+/**
+ * A hook that mirrors `useState` in function and API, expect that setState
+ * calls return a promise that resolves after the state has been set (in an effect).
+ *
+ * This is _similar_ to the second callback in classy setState calls, but fires later.
+ *
+ * ```ts
+ * const [counter, setState] = useStateAsync(1);
+ *
+ * const handleIncrement = async () => {
+ *   await setState(2);
+ *   doWorkRequiringCurrentState()
+ * }
+ * ```
+ *
+ * @param initialState initialize with some state value same as `useState`
+ */
+function useStateAsync<TState>(
+  initialState: TState | (() => TState),
+): [TState, AsyncSetState<TState>] {
+  const [state, setState] = useState(initialState)
+  const resolvers = useRef<((state: TState) => void)[]>([])
+
+  useEffect(() => {
+    resolvers.current.forEach(resolve => resolve(state))
+    resolvers.current.length = 0
+  }, [state])
+
+  const setStateAsync = useCallback(
+    (update: Updater<TState> | TState) => {
+      return new Promise<TState>((resolve, reject) => {
+        setState(prevState => {
+          try {
+            let nextState: TState
+            // ugly instanceof for typescript
+            if (update instanceof Function) {
+              nextState = update(prevState)
+            } else {
+              nextState = update
+            }
+
+            // If state does not change, we must resolve the promise because
+            // react won't re-render and effect will not resolve. If there are already
+            // resolvers queued, then it should be safe to assume an update will happen
+            if (!resolvers.current.length && Object.is(nextState, prevState)) {
+              resolve(nextState)
+            } else {
+              resolvers.current.push(resolve)
+            }
+            return nextState
+          } catch (e) {
+            reject(e)
+            throw e
+          }
+        })
+      })
+    },
+    [setState],
+  )
+  return [state, setStateAsync]
+}
+
+export default useStateAsync

test/setup.js

@@ -1,5 +1,5 @@
-import Adapter from 'enzyme-adapter-react-16'
 import Enzyme from 'enzyme'
+import Adapter from 'enzyme-adapter-react-16'
 import matchMediaPolyfill from 'mq-polyfill'
 
 Enzyme.configure({ adapter: new Adapter() })
@@ -22,3 +22,35 @@ if (typeof window !== 'undefined') {
     }).dispatchEvent(new this.Event('resize'))
   }
 }
+
+let expectedErrors = 0
+let actualErrors = 0
+function onError(e) {
+  if (expectedErrors) {
+    e.preventDefault()
+  }
+  actualErrors += 1
+}
+
+expect.errors = num => {
+  expectedErrors = num
+}
+
+beforeEach(() => {
+  expectedErrors = 0
+  actualErrors = 0
+  if (typeof window !== 'undefined') {
+    window.addEventListener('error', onError)
+  }
+})
+
+afterEach(() => {
+  if (typeof window !== 'undefined') {
+    window.removeEventListener('error', onError)
+  }
+  if (expectedErrors) {
+    expect(actualErrors).toBe(expectedErrors)
+  }
+
+  expectedErrors = 0
+})

test/useStateAsync.test.tsx

@@ -0,0 +1,109 @@
+import { mount } from 'enzyme'
+import React from 'react'
+import { act } from 'react-dom/test-utils'
+import useStateAsync, { AsyncSetState } from '../src/useStateAsync'
+
+describe('useStateAsync', () => {
+  it('should increment counter', async () => {
+    let asyncState: [number, AsyncSetState<number>]
+
+    function Wrapper() {
+      asyncState = useStateAsync<number>(0)
+      return null
+    }
+
+    mount(<Wrapper />)
+
+    expect.assertions(4)
+
+    const incrementAsync = async () => {
+      await act(() => asyncState[1](prev => prev + 1))
+    }
+
+    expect(asyncState![0]).toEqual(0)
+
+    await incrementAsync()
+    expect(asyncState![0]).toEqual(1)
+
+    await incrementAsync()
+    expect(asyncState![0]).toEqual(2)
+
+    await incrementAsync()
+    expect(asyncState![0]).toEqual(3)
+  })
+
+  it('should reject on error', async () => {
+    let asyncState: [number, AsyncSetState<number>]
+
+    function Wrapper() {
+      asyncState = useStateAsync<number>(1)
+      return null
+    }
+    class CatchError extends React.Component {
+      static getDerivedStateFromError() {}
+      componentDidCatch() {}
+      render() {
+        return this.props.children
+      }
+    }
+
+    mount(
+      <CatchError>
+        <Wrapper />
+      </CatchError>,
+    )
+
+    // @ts-ignore
+    expect.errors(1)
+
+    await act(async () => {
+      const p = asyncState[1](() => {
+        throw new Error('yo')
+      })
+      return expect(p).rejects.toThrow('yo')
+    })
+  })
+
+  it('should resolve even if no update happens', async () => {
+    let asyncState: [number, AsyncSetState<number>]
+
+    function Wrapper() {
+      asyncState = useStateAsync<number>(1)
+      return null
+    }
+
+    mount(<Wrapper />)
+
+    expect.assertions(3)
+
+    expect(asyncState![0]).toEqual(1)
+
+    await act(() => expect(asyncState[1](1)).resolves.toEqual(1))
+
+    expect(asyncState![0]).toEqual(1)
+  })
+
+  it('should resolve after update if already pending', async () => {
+    let asyncState: [number, AsyncSetState<number>]
+
+    function Wrapper() {
+      asyncState = useStateAsync<number>(0)
+      return null
+    }
+
+    mount(<Wrapper />)
+
+    expect.assertions(5)
+
+    expect(asyncState![0]).toEqual(0)
+
+    const setAndAssert = async (n: number) =>
+      expect(asyncState[1](n)).resolves.toEqual(2)
+
+    await act(() =>
+      Promise.all([setAndAssert(1), setAndAssert(1), setAndAssert(2)]),
+    )
+
+    expect(asyncState![0]).toEqual(2)
+  })
+})