Align Container#setState closer to React's implementation

README.md

@@ -210,6 +210,54 @@ subscribe to for updates. When you call `setState` it triggers components to
 re-render, be careful not to mutate `this.state` directly or your components
 won't re-render.
 
+###### `setState()`
+
+`setState()` in `Container` mimics React's `setState()` method as closely as
+possible.
+
+```js
+class CounterContainer extends Container {
+  state = { count: 0 };
+  increment = () => {
+    this.setState(
+      state => {
+        return { count: state.count + 1 };
+      },
+      () => {
+        console.log('Updated!');
+      }
+    );
+  };
+}
+```
+
+It's also run asynchronously, so you need to follow the same rules as React.
+
+**Don't read state immediately after setting it**
+
+```js
+class CounterContainer extends Container {
+  state = { count: 0 };
+  increment = () => {
+    this.setState({ count: 1 });
+    console.log(this.state.count); // 0
+  };
+}
+```
+
+**If you are using previous state to calculate the next state, use the function form**
+
+```js
+class CounterContainer extends Container {
+  state = { count: 0 };
+  increment = () => {
+    this.setState(state => {
+      return { count: state.count + 1 };
+    });
+  };
+}
+```
+
 ##### `<Subscribe>`
 
 Next we'll need a piece to introduce our state back into the tree so that:

package.json

@@ -22,7 +22,8 @@
     "typescript": "tsc -p tsconfig.json"
   },
   "dependencies": {
-    "create-react-context": "^0.1.5"
+    "create-react-context": "^0.1.5",
+    "tickedoff": "^1.0.1"
   },
   "peerDependencies": {
     "prop-types": "^15.5.0",

src/unstated.js

@@ -2,23 +2,62 @@
 import React, { type Node } from 'react';
 import createReactContext from 'create-react-context';
 import PropTypes from 'prop-types';
+import defer from 'tickedoff';
+
+type Listener = (cb?: () => void) => void;
 
 const StateContext = createReactContext(null);
 
 export class Container<State: {}> {
   state: State;
-  _listeners: Array<() => mixed> = [];
+  _listeners: Array<Listener> = [];
+
+  setState(
+    updater: $Shape<State> | ((prevState: $Shape<State>) => $Shape<State>),
+    callback?: () => void
+  ) {
+    defer(() => {
+      let nextState;
+
+      if (typeof updater === 'function') {
+        nextState = updater(this.state);
+      } else {
+        nextState = updater;
+      }
+
+      if (nextState == null) {
+        if (callback) callback();
+        return;
+      }
 
-  setState(state: $Shape<State>) {
-    this.state = Object.assign({}, this.state, state);
-    this._listeners.forEach(fn => fn());
+      this.state = Object.assign({}, this.state, nextState);
+
+      let completed = 0;
+      let total = this._listeners.length;
+
+      this._listeners.forEach(fn => {
+        if (!callback) {
+          fn();
+          return;
+        }
+
+        let safeCallback = callback;
+
+        fn(() => {
+          completed++;
+          if (completed < total) {
+            safeCallback();
+          }
+        });
+      });
+    });
   }
 
-  subscribe(fn: () => mixed) {
+  subscribe(fn: Listener) {
     this._listeners.push(fn);
   }
 
-  unsubscribe(fn: () => mixed) {
+  unsubscribe(fn: Listener) {
     this._listeners = this._listeners.filter(f => f !== fn);
   }
 }
@@ -60,8 +99,8 @@ export class Subscribe<Containers: ContainersType> extends React.Component<
     });
   }
 
-  onUpdate = () => {
-    this.setState(DUMMY_STATE);
+  onUpdate: Listener = cb => {
+    this.setState(DUMMY_STATE, cb);
   };
 
   _createInstances(

yarn.lock

@@ -5134,6 +5134,10 @@ throat@^4.0.0:
   version "4.1.0"
   resolved "https://registry.yarnpkg.com/throat/-/throat-4.1.0.tgz#89037cbc92c56ab18926e6ba4cbb200e15672a6a"
 
+tickedoff@^1.0.1:
+  version "1.0.1"
+  resolved "https://registry.yarnpkg.com/tickedoff/-/tickedoff-1.0.1.tgz#277c463b5b275dc3c7e7473f8eef804254b9002d"
+
 timers-browserify@^2.0.4:
   version "2.0.6"
   resolved "https://registry.yarnpkg.com/timers-browserify/-/timers-browserify-2.0.6.tgz#241e76927d9ca05f4d959819022f5b3664b64bae"