Polyfill for getSnapshotBeforeUpdate

index.js

@@ -26,10 +26,24 @@ function componentWillReceiveProps(nextProps) {
   }
 }
 
+function componentWillUpdate(nextProps, nextState) {
+  var prevProps = this.props;
+  var prevState = this.state;
+  this.props = nextProps;
+  this.state = nextState;
+  this.__reactInternalSnapshot = this.getSnapshotBeforeUpdate(
+    prevProps,
+    prevState
+  );
+  this.props = prevProps;
+  this.state = prevState;
+}
+
 // React may warn about cWM/cWRP/cWU methods being deprecated.
 // Add a flag to suppress these warnings for this special case.
 componentWillMount.__suppressDeprecationWarning = true;
 componentWillReceiveProps.__suppressDeprecationWarning = true;
+componentWillUpdate.__suppressDeprecationWarning = true;
 
 module.exports = function polyfill(Component) {
   if (!Component.prototype || !Component.prototype.isReactComponent) {
@@ -51,5 +65,33 @@ module.exports = function polyfill(Component) {
     Component.prototype.componentWillReceiveProps = componentWillReceiveProps;
   }
 
+  if (typeof Component.prototype.getSnapshotBeforeUpdate === 'function') {
+    if (typeof Component.prototype.componentWillUpdate === 'function') {
+      throw new Error('Cannot polyfill if componentWillUpdate already exists');
+    }
+
+    if (typeof Component.prototype.componentDidUpdate !== 'function') {
+      throw new Error(
+        'Cannot polyfill getSnapshotBeforeUpdate() unless componentDidUpdate() exists on the prototype'
+      );
+    }
+
+    Component.prototype.componentWillUpdate = componentWillUpdate;
+
+    var componentDidUpdate = Component.prototype.componentDidUpdate;
+
+    Component.prototype.componentDidUpdate = function componentDidUpdatePolyfill(
+      prevProps,
+      prevState
+    ) {
+      componentDidUpdate.call(
+        this,
+        prevProps,
+        prevState,
+        this.__reactInternalSnapshot
+      );
+    };
+  }
+
   return Component;
 };

test.js

@@ -54,6 +54,8 @@ Object.entries(POLYFILLS).forEach(([name, polyfill]) => {
         });
 
         it('should support create-react-class components', () => {
+          let componentDidUpdateCalled = false;
+
           const CRCComponent = createReactClass({
             statics: {
               getDerivedStateFromProps(nextProps, prevState) {
@@ -65,6 +67,17 @@ Object.entries(POLYFILLS).forEach(([name, polyfill]) => {
             getInitialState() {
               return {count: 1};
             },
+            getSnapshotBeforeUpdate(prevProps, prevState) {
+              return prevState.count * 2 + this.state.count * 3;
+            },
+            componentDidUpdate(prevProps, prevState, snapshot) {
+              expect(prevProps).toEqual({incrementBy: 2});
+              expect(prevState).toEqual({count: 3});
+              expect(this.props).toEqual({incrementBy: 3});
+              expect(this.state).toEqual({count: 6});
+              expect(snapshot).toBe(24);
+              componentDidUpdateCalled = true;
+            },
             render() {
               return React.createElement('div', null, this.state.count);
             },
@@ -79,13 +92,86 @@ Object.entries(POLYFILLS).forEach(([name, polyfill]) => {
           );
 
           expect(container.textContent).toBe('3');
+          expect(componentDidUpdateCalled).toBe(false);
 
           ReactDOM.render(
             React.createElement(CRCComponent, {incrementBy: 3}),
             container
           );
 
           expect(container.textContent).toBe('6');
+          expect(componentDidUpdateCalled).toBe(true);
+        });
+
+        it('should support class components', () => {
+          let componentDidUpdateCalled = false;
+
+          class Component extends React.Component {
+            constructor(props) {
+              super(props);
+              this.state = {count: 1};
+            }
+            static getDerivedStateFromProps(nextProps, prevState) {
+              return {
+                count: prevState.count + nextProps.incrementBy,
+              };
+            }
+            getSnapshotBeforeUpdate(prevProps, prevState) {
+              return prevState.count * 2 + this.state.count * 3;
+            }
+            componentDidUpdate(prevProps, prevState, snapshot) {
+              expect(prevProps).toEqual({incrementBy: 2});
+              expect(prevState).toEqual({count: 3});
+              expect(this.props).toEqual({incrementBy: 3});
+              expect(this.state).toEqual({count: 6});
+              expect(snapshot).toBe(24);
+              componentDidUpdateCalled = true;
+            }
+            render() {
+              return React.createElement('div', null, this.state.count);
+            }
+          }
+
+          polyfill(Component);
+
+          const container = document.createElement('div');
+          ReactDOM.render(
+            React.createElement(Component, {incrementBy: 2}),
+            container
+          );
+
+          expect(container.textContent).toBe('3');
+          expect(componentDidUpdateCalled).toBe(false);
+
+          ReactDOM.render(
+            React.createElement(Component, {incrementBy: 3}),
+            container
+          );
+
+          expect(container.textContent).toBe('6');
+          expect(componentDidUpdateCalled).toBe(true);
+        });
+
+        it('should throw if componentDidUpdate is not defined on the prototype', () => {
+          let componentDidUpdateCalled = false;
+
+          class Component extends React.Component {
+            constructor(props) {
+              super(props);
+
+              Object.defineProperty(this, 'componentDidUpdate', {
+                value: (prevProps, prevState, snapshot) => {},
+              });
+            }
+            getSnapshotBeforeUpdate(prevProps, prevState) {}
+            render() {
+              return null;
+            }
+          }
+
+          expect(() => polyfill(Component)).toThrow(
+            'Cannot polyfill getSnapshotBeforeUpdate() unless componentDidUpdate() exists on the prototype'
+          );
         });
 
         it('should support getDerivedStateFromProps in subclass', () => {
@@ -174,6 +260,20 @@ Object.entries(POLYFILLS).forEach(([name, polyfill]) => {
             'Cannot polyfill if componentWillReceiveProps already exists'
           );
         });
+
+        it('should error if component already has cWU lifecycles with gSBU', () => {
+          class ComponentWithWillUpdate extends React.Component {
+            componentWillUpdate() {}
+            getSnapshotBeforeUpdate() {}
+            render() {
+              return null;
+            }
+          }
+
+          expect(() => polyfill(ComponentWithWillUpdate)).toThrow(
+            'Cannot polyfill if componentWillUpdate already exists'
+          );
+        });
       });
     });
   });