Avoid work inside of render()

Our HOC here is currently doing more work than is really necessary
inside of render. This can make updates slower than they need to be. To
optimize this some, I'm moving this work into the constructor and
componentWillReceiveProps to avoid unnecessarily repeating it.

src/withStyles.jsx

@@ -37,14 +37,37 @@ export function withStyles(
   const styleDef = styleFn && ThemedStyleSheet.create(styleFn);
   const BaseClass = baseClass(pureComponent);
 
+  function getAddedProps(themeName) {
+    const addedProps = {
+      [themePropName]: ThemedStyleSheet.get(themeName),
+    };
+
+    if (styleDef) {
+      addedProps[stylesPropName] = styleDef(themeName);
+    }
+
+    return addedProps;
+  }
+
   return function withStylesHOC(WrappedComponent) {
-    // NOTE: Use a class here so components are ref-able if need be:
-    // eslint-disable-next-line react/prefer-stateless-function
     class WithStyles extends BaseClass {
-      render() {
-        const props = this.props;
-        const { themeName } = this.context;
+      constructor(props, context) {
+        super(props, context);
+
+        this.state = {
+          addedProps: getAddedProps(context.themeName),
+        };
+      }
 
+      componentWillReceiveProps(nextProps, nextContext) {
+        if (this.context.themeName !== nextContext.themeName) {
+          this.setState({
+            addedProps: getAddedProps(nextContext.themeName),
+          });
+        }
+      }
+
+      render() {
         // As some components will depend on previous styles in
         // the component tree, we provide the option of flushing the
         // buffered styles (i.e. to a style tag) **before** the rendering
@@ -56,15 +79,7 @@ export function withStyles(
           ThemedStyleSheet.flush();
         }
 
-        const addedProps = {
-          [themePropName]: ThemedStyleSheet.get(themeName),
-        };
-
-        if (styleDef) {
-          addedProps[stylesPropName] = styleDef(themeName);
-        }
-
-        return <WrappedComponent {...props} {...addedProps} />;
+        return <WrappedComponent {...this.props} {...this.state.addedProps} />;
       }
     }
 

test/withStyles_test.jsx

@@ -120,6 +120,36 @@ describe('withStyles()', () => {
       shallow(<Wrapped />, { context: { themeName: 'tropical' } });
     });
 
+    it('reacts to changed context', () => {
+      const tropicalTheme = {
+        color: {
+          red: 'yellow',
+        },
+      };
+      ThemedStyleSheet.registerTheme('tropical', tropicalTheme);
+
+      function MyComponent() {
+        return null;
+      }
+
+      const Wrapped = withStyles(({ color }) => ({
+        foo: {
+          color: color.red,
+        },
+      }))(MyComponent);
+      const wrapper = shallow(<Wrapped />, { context: { themeName: 'default' } });
+
+      // default theme
+      expect(wrapper.state('addedProps').styles)
+        .to.eql({ foo: { color: '#990000' } });
+
+      wrapper.setContext({ themeName: 'tropical' });
+
+      // tropical theme
+      expect(wrapper.state('addedProps').styles)
+        .to.eql({ foo: { color: 'yellow' } });
+    });
+
     it('allows the styles prop name to be customized', () => {
       function MyComponent({ bar }) {
         expect(bar).to.eql({ foo: { color: '#ff0000' } });