fix: use React.forwardRef and improve memoization

src/__tests__/createTheming.test.js

@@ -1,5 +1,6 @@
 /* eslint-disable react/no-multi-comp */
-import React, { Fragment } from 'react';
+
+import * as React from 'react';
 import ReactDOM from 'react-dom';
 import createTheming from '../createTheming';
 
@@ -26,8 +27,6 @@ describe('createTheming', () => {
     secondaryColor: '#ffffff',
   };
 
-  const originalTheme = { ...lightTheme };
-
   const { ThemeProvider, withTheme, useTheme } = createTheming(darkTheme);
 
   it('provides theme prop with HOC', () => {
@@ -115,7 +114,7 @@ describe('createTheming', () => {
     );
   });
 
-  it('allows to use ref on wrapped component', () => {
+  it('sets correct ref on wrapped component', () => {
     class Component extends React.Component {
       foo() {
         return 'bar';
@@ -125,12 +124,11 @@ describe('createTheming', () => {
         return null;
       }
     }
+
     const WithThemeComponent = withTheme(Component);
 
     class Wrapper extends React.Component {
       componentDidMount() {
-        expect(typeof this.comp).toBe('object');
-        expect(typeof this.comp.foo).toEqual('function');
         expect(this.comp.foo()).toEqual('bar');
       }
 
@@ -153,71 +151,6 @@ describe('createTheming', () => {
     );
   });
 
-  it('Should set properly getWrappedInstance method', () => {
-    class Component extends React.Component {
-      foo() {
-        return 'bar';
-      }
-
-      getWrappedInstance() {
-        return this;
-      }
-
-      render() {
-        return null;
-      }
-    }
-    class ComponentWithoutGWI extends React.Component {
-      foo() {
-        return 'bar';
-      }
-
-      render() {
-        return null;
-      }
-    }
-    const WithThemeComponent = withTheme(Component);
-    const WithThemeComponentWithoutGWI = withTheme(ComponentWithoutGWI);
-
-    class Wrapper extends React.Component {
-      componentDidMount() {
-        expect(typeof this.comp).toBe('object');
-        expect(typeof this.comp.getWrappedInstance).toEqual('function');
-        expect(this.comp.getWrappedInstance().foo()).toEqual('bar');
-
-        expect(typeof this.compWithoutGWI).toBe('object');
-        expect(typeof this.compWithoutGWI.getWrappedInstance).toEqual(
-          'function'
-        );
-        expect(this.compWithoutGWI.getWrappedInstance().foo()).toEqual('bar');
-      }
-
-      render() {
-        return (
-          <Fragment>
-            <WithThemeComponent
-              ref={component => {
-                this.comp = component;
-              }}
-            />
-            <WithThemeComponentWithoutGWI
-              ref={component => {
-                this.compWithoutGWI = component;
-              }}
-            />
-          </Fragment>
-        );
-      }
-    }
-
-    ReactDOM.render(
-      <ThemeProvider>
-        <Wrapper />
-      </ThemeProvider>,
-      node
-    );
-  });
-
   it('merge theme from provider and prop', () => {
     const PropsChecker = withTheme(({ theme }) => {
       expect(theme).not.toBe(lightTheme);
@@ -290,26 +223,25 @@ describe('createTheming', () => {
   });
 
   it('doesnt mutate existing theme', () => {
+    const overrides = { primaryColor: 'red' };
     const Checker1 = withTheme(({ theme }) => {
-      expect(theme).toEqual({
-        ...lightTheme,
-        primaryColor: 'red',
-      });
+      expect(theme).not.toBe(lightTheme);
+      expect(theme).not.toBe(overrides);
       return null;
     });
 
     const Checker1WithTheme = withTheme(Checker1);
 
     const Checker2 = withTheme(({ theme }) => {
-      expect(theme).toEqual(originalTheme);
+      expect(theme).toBe(lightTheme);
       return null;
     });
 
     const Checker2WithTheme = withTheme(Checker2);
 
     ReactDOM.render(
       <ThemeProvider theme={lightTheme}>
-        <Checker1WithTheme theme={{ primaryColor: 'red' }} />
+        <Checker1WithTheme theme={overrides} />
         <Checker2WithTheme />
       </ThemeProvider>,
       node

src/createWithTheme.js

@@ -4,14 +4,9 @@ import * as React from 'react';
 import deepmerge from 'deepmerge';
 import hoistNonReactStatics from 'hoist-non-react-statics';
 
-import { copyRefs } from './utils';
-
 import type { ThemeProviderType } from './createThemeProvider';
 import type { $DeepShape } from './types';
 
-const isClassComponent = (Component: any) =>
-  Boolean(Component.prototype && Component.prototype.isReactComponent);
-
 export type WithThemeType<T> = <P, C: React.ComponentType<P>>(
   Comp: C
 ) => C &
@@ -25,8 +20,6 @@ const createWithTheme = <T: Object, S: $DeepShape<T>>(
 ) =>
   function withTheme(Comp: *) {
     class ThemedComponent extends React.Component<*> {
-      static displayName = `withTheme(${Comp.displayName || Comp.name})`;
-
       _previous: ?{ a: T, b: ?S, result: T };
 
       _merge = (a: T, b: ?S) => {
@@ -36,39 +29,29 @@ const createWithTheme = <T: Object, S: $DeepShape<T>>(
           return previous.result;
         }
 
-        const result = a && b ? deepmerge(a, b) : a || b;
+        const result = a && b && a !== b ? deepmerge(a, b) : a || b;
 
         this._previous = { a, b, result };
 
         return result;
       };
 
-      _root: any;
-
       render() {
+        const { _reactThemeProviderForwardedRef, ...rest } = this.props;
+
         return (
           <ThemeContext.Consumer>
             {theme => {
-              const merged = this._merge(theme, this.props.theme);
-
-              let element;
-              if (isClassComponent(Comp)) {
-                // Only add refs for class components as function components don't support them
-                // It's needed to support use cases which need access to the underlying node
-                element = (
-                  <Comp
-                    {...this.props}
-                    ref={c => {
-                      this._root = c;
-                    }}
-                    theme={merged}
-                  />
-                );
-              } else {
-                element = <Comp {...this.props} theme={merged} />;
-              }
-
-              if (merged !== this.props.theme) {
+              const merged = this._merge(theme, rest.theme);
+              const element = (
+                <Comp
+                  {...rest}
+                  theme={merged}
+                  ref={_reactThemeProviderForwardedRef}
+                />
+              );
+
+              if (rest.theme && merged !== rest.theme) {
                 // If a theme prop was passed, expose it to the children
                 return <ThemeProvider theme={merged}>{element}</ThemeProvider>;
               }
@@ -80,23 +63,15 @@ const createWithTheme = <T: Object, S: $DeepShape<T>>(
       }
     }
 
-    if (isClassComponent(Comp)) {
-      // getWrappedInstance is exposed by some HOCs like react-redux's connect
-      // Use it to get the ref to the underlying element
-      // Also expose it to access the underlying element after wrapping
-      // $FlowFixMe
-      ThemedComponent.prototype.getWrappedInstance = function getWrappedInstance() {
-        return this._root && this._root.getWrappedInstance
-          ? this._root.getWrappedInstance()
-          : this._root;
-      };
+    const ResultComponent = React.forwardRef((props, ref) => (
+      <ThemedComponent {...props} _reactThemeProviderForwardedRef={ref} />
+    ));
 
-      ThemedComponent = copyRefs(ThemedComponent, Comp);
-    }
+    ResultComponent.displayName = `withTheme(${Comp.displayName || Comp.name})`;
 
-    hoistNonReactStatics(ThemedComponent, Comp);
+    hoistNonReactStatics(ResultComponent, Comp);
 
-    return (ThemedComponent: any);
+    return (ResultComponent: any);
   };
 
 export default createWithTheme;