Recalculate state and bindings on hot reload. Fixes #27

Author: gaearon

src/components/createConnect.js

@@ -30,6 +30,9 @@ function areStatePropsEqual(stateProps, nextStateProps) {
   return shallowEqual(stateProps, nextStateProps);
 }
 
+// Helps track hot reloading.
+let nextVersion = 0;
+
 export default function createConnect(React) {
   const { Component, PropTypes } = React;
   const storeShape = createStoreShape(PropTypes);
@@ -44,6 +47,9 @@ export default function createConnect(React) {
       wrapActionCreators(actionCreatorsOrMapDispatchToProps) :
       actionCreatorsOrMapDispatchToProps;
 
+    // Helps track hot reloading.
+    const version = nextVersion++;
+
     return DecoratedComponent => class Connect extends Component {
       static displayName = `Connect(${getDisplayName(DecoratedComponent)})`;
       static DecoratedComponent = DecoratedComponent;
@@ -61,6 +67,7 @@ export default function createConnect(React) {
 
       constructor(props, context) {
         super(props, context);
+        this.version = version;
         this.setUnderlyingRef = ::this.setUnderlyingRef;
         this.state = {
           ...this.mapState(props, context),
@@ -72,18 +79,45 @@ export default function createConnect(React) {
         return typeof this.unsubscribe === 'function';
       }
 
-      componentDidMount() {
-        if (shouldSubscribe) {
+      trySubscribe() {
+        if (shouldSubscribe && !this.unsubscribe) {
           this.unsubscribe = this.context.store.subscribe(::this.handleChange);
         }
       }
 
-      componentWillUnmount() {
+      tryUnsubscribe() {
         if (this.isSubscribed()) {
           this.unsubscribe();
+          this.unsubscribe = null;
         }
       }
 
+      componentDidMount() {
+        this.trySubscribe();
+      }
+
+      componentWillUpdate() {
+        if (process.env.NODE_ENV !== 'production') {
+          if (this.version === version) {
+            return;
+          }
+
+          // We are hot reloading!
+          this.version = version;
+
+          // Update the state and bindings.
+          this.trySubscribe();
+          this.setState({
+            ...this.mapState(),
+            ...this.mapDispatch()
+          });
+        }
+      }
+
+      componentWillUnmount() {
+        this.tryUnsubscribe();
+      }
+
       handleChange(props = this.props) {
         const nextState = this.mapState(props, this.context);
         if (!areStatePropsEqual(this.state.stateProps, nextState.stateProps)) {

test/components/connect.spec.js

@@ -435,6 +435,57 @@ describe('React', () => {
       }).toThrow(/mergeProps/);
     });
 
+    it('should recalculate the state and rebind the actions on hot update', () => {
+      const store = createStore(() => {});
+
+      @connect(
+        undefined,
+        () => ({ scooby: 'doo' })
+      )
+      class ContainerBefore extends Component {
+        render() {
+          return (
+              <div {...this.props} />
+          );
+        }
+      }
+
+      @connect(
+        () => ({ foo: 'baz' }),
+        () => ({ scooby: 'foo' })
+      )
+      class ContainerAfter extends Component {
+        render() {
+          return (
+              <div {...this.props} />
+          );
+        }
+      }
+
+      let container;
+      TestUtils.renderIntoDocument(
+        <Provider store={store}>
+          {() => <ContainerBefore ref={instance => container = instance} />}
+         </Provider>
+      );
+      const div = TestUtils.findRenderedDOMComponentWithTag(container, 'div');
+      expect(div.props.foo).toEqual(undefined);
+      expect(div.props.scooby).toEqual('doo');
+
+      // Crude imitation of hot reloading that does the job
+      Object.keys(ContainerAfter.prototype).filter(key =>
+        typeof ContainerAfter.prototype[key] === 'function'
+      ).forEach(key => {
+        if (key !== 'render') {
+          ContainerBefore.prototype[key] = ContainerAfter.prototype[key];
+        }
+      });
+
+      container.forceUpdate();
+      expect(div.props.foo).toEqual('baz');
+      expect(div.props.scooby).toEqual('foo');
+    });
+
     it('should set the displayName correctly', () => {
       expect(connect(state => state)(
         class Foo extends Component {