wip

Author: jamiebuilds

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: Function) {
+  subscribe(fn: Listener) {
     this._listeners.push(fn);
   }
 
-  unsubscribe(fn: Function) {
+  unsubscribe(fn: Listener) {
     this._listeners = this._listeners.filter(f => f !== fn);
   }
 }
@@ -60,16 +99,16 @@ export class Subscribe<Containers: ContainersType> extends React.Component<
     });
   }
 
-  onUpdate = () => {
-    this.setState(DUMMY_STATE);
+  onUpdate: Listener = cb => {
+    this.setState(DUMMY_STATE, cb);
   };
 
   _createInstances(
     map: ContainerMapType | null,
     containers: ContainersType
   ): Array<ContainerType> {
     this._unsubscribe();
-                       
+
     if (map === null) {
       throw new Error(
         'You must wrap your <Subscribe> components with a <Provider>'

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"