From 8105d8fd512cb23110ff81508157dbfe51b165e2 Mon Sep 17 00:00:00 2001
From: Tica Costica Puntaru <ticaasd@gmail.com>
Date: Tue, 25 Sep 2018 17:58:24 +0300
Subject: [PATCH] Update index.js

implement setState done callback
---
 src/index.js | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/index.js b/src/index.js
index cda36e9..8f54493 100644
--- a/src/index.js
+++ b/src/index.js
@@ -216,7 +216,7 @@ export default function reactTreeWalker(
               instance.state = instance.state || null
 
               // Make the setState synchronous.
-              instance.setState = newState => {
+              instance.setState = (newState, cb) => {
                 if (typeof newState === 'function') {
                   // eslint-disable-next-line no-param-reassign
                   newState = newState(
@@ -226,6 +226,9 @@ export default function reactTreeWalker(
                   )
                 }
                 instance.state = Object.assign({}, instance.state, newState)
+                if (typeof cb === 'function'){
+                  cb();
+                }
               }
 
               if (Component.getDerivedStateFromProps) {