Support updater function in setState

https://reactjs.org/docs/react-component.html#setstate Based on https://github.com/apollographql/react-apollo/blob/a42f7da5bd4163223a87ab545dceb88e778046fe/src/getDataFromTree.ts#L105
ctrlplusb · Feb 13, 2018 · 4a898bf · 4a898bf
1 parent 5951226
commit 4a898bf
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/src/__tests__/index.test.js b/src/__tests__/index.test.js
@@ -101,6 +101,7 @@ describe('reactTreeWalker', () => {
 
       componentWillMount() {
         this.setState({ foo: 'bar' })
+        this.setState((state, props) => ({ other: `I am ${props.value} ${state.foo}` }))
       }
 
       render() {
@@ -109,8 +110,8 @@ describe('reactTreeWalker', () => {
       }
     }
 
-    return reactTreeWalker(<Baz />, () => true).then(() => {
-      const expected = { foo: 'bar' }
+    return reactTreeWalker(<Baz value="foo" />, () => true).then(() => {
+      const expected = { foo: 'bar', other: 'I am foo bar' }
       expect(actual).toMatchObject(expected)
     })
   })

diff --git a/src/index.js b/src/index.js
@@ -136,6 +136,14 @@ export default function reactTreeWalker(
 
         // Make the setState synchronous.
         instance.setState = newState => {
+          if (typeof newState === 'function') {
+            // eslint-disable-next-line no-param-reassign
+            newState = newState(
+              instance.state,
+              instance.props,
+              instance.context,
+            )
+          }
           instance.state = Object.assign({}, instance.state, newState)
         }