[HOC] Add innerRef to withWidth and withTheme (#12236)

packages/material-ui/src/styles/withTheme.d.ts

@@ -3,6 +3,7 @@ import { ConsistentWith } from '..';
 
 export interface WithTheme {
   theme: Theme;
+  innerRef?: React.Ref<any> | React.RefObject<any>;
 }
 
 declare const withTheme: () => <P extends ConsistentWith<P, WithTheme>>(

packages/material-ui/src/styles/withTheme.js

@@ -1,4 +1,5 @@
 import React from 'react';
+import PropTypes from 'prop-types';
 import hoistNonReactStatics from 'hoist-non-react-statics';
 import wrapDisplayName from 'recompose/wrapDisplayName';
 import createMuiTheme from './createMuiTheme';
@@ -44,10 +45,18 @@ const withTheme = () => Component => {
     }
 
     render() {
-      return <Component theme={this.state.theme} {...this.props} />;
+      const { innerRef, ...other } = this.props;
+      return <Component theme={this.state.theme} ref={innerRef} {...other} />;
     }
   }
 
+  WithTheme.propTypes = {
+    /**
+     * Use that property to pass a ref callback to the decorated component.
+     */
+    innerRef: PropTypes.oneOfType([PropTypes.func, PropTypes.object]),
+  };
+
   WithTheme.contextTypes = themeListener.contextTypes;
 
   if (process.env.NODE_ENV !== 'production') {

packages/material-ui/src/styles/withTheme.test.js

@@ -1,11 +1,18 @@
 import React from 'react';
 import { assert } from 'chai';
+import { spy } from 'sinon';
 import createBroadcast from 'brcast';
 import { createShallow, createMount } from '../test-utils';
 import { CHANNEL } from './themeListener';
 import withTheme from './withTheme';
 
 const Empty = () => <div />;
+// eslint-disable-next-line react/prefer-stateless-function
+class EmptyClass extends React.Component<{}> {
+  render() {
+    return <div />;
+  }
+}
 
 describe('withTheme', () => {
   let shallow;
@@ -44,4 +51,13 @@ describe('withTheme', () => {
     broadcast.setState(newTheme);
     assert.strictEqual(wrapper.instance().state.theme, newTheme);
   });
+
+  describe('prop: innerRef', () => {
+    it('should provide a ref on the inner component', () => {
+      const ThemedComponent = withTheme()(EmptyClass);
+      const handleRef = spy();
+      mount(<ThemedComponent innerRef={handleRef} />);
+      assert.strictEqual(handleRef.callCount, 1);
+    });
+  });
 });

packages/material-ui/src/withWidth/withWidth.d.ts

@@ -7,6 +7,7 @@ export interface WithWidthOptions {
 
 export interface WithWidthProps {
   width: Breakpoint;
+  innerRef?: React.Ref<any> | React.RefObject<any>;
 }
 
 export function isWidthDown(

packages/material-ui/src/withWidth/withWidth.js

@@ -98,7 +98,7 @@ const withWidth = (options = {}) => Component => {
     }
 
     render() {
-      const { initialWidth, theme, width, ...other } = this.props;
+      const { initialWidth, theme, width, innerRef, ...other } = this.props;
 
       const props = {
         width:
@@ -127,7 +127,7 @@ const withWidth = (options = {}) => Component => {
 
       return (
         <EventListener target="window" onResize={this.handleResize}>
-          <Component {...more} {...props} />
+          <Component {...more} {...props} ref={innerRef} />
         </EventListener>
       );
     }
@@ -144,6 +144,10 @@ const withWidth = (options = {}) => Component => {
      * http://caniuse.com/#search=client%20hint
      */
     initialWidth: PropTypes.oneOf(['xs', 'sm', 'md', 'lg', 'xl']),
+    /**
+     * Use that property to pass a ref callback to the decorated component.
+     */
+    innerRef: PropTypes.oneOfType([PropTypes.func, PropTypes.object]),
     /**
      * @ignore
      */

packages/material-ui/src/withWidth/withWidth.test.js

@@ -1,12 +1,19 @@
 import React from 'react';
 import { assert } from 'chai';
-import { useFakeTimers } from 'sinon';
+import { useFakeTimers, spy } from 'sinon';
 import { createMount, createShallow } from '../test-utils';
 import withWidth, { isWidthDown, isWidthUp } from './withWidth';
 import createBreakpoints from '../styles/createBreakpoints';
 import createMuiTheme from '../styles/createMuiTheme';
 
 const Empty = () => <div />;
+// eslint-disable-next-line react/prefer-stateless-function
+class EmptyClass extends React.Component<{}> {
+  render() {
+    return <div />;
+  }
+}
+const EmptyClassWithWidth = withWidth()(EmptyClass);
 const EmptyWithWidth = withWidth()(Empty);
 
 const breakpoints = createBreakpoints({});
@@ -39,6 +46,15 @@ describe('withWidth', () => {
     });
   });
 
+  describe('prop: innerRef', () => {
+    it('should provide a ref on the inner component', () => {
+      const handleRef = spy();
+
+      mount(<EmptyClassWithWidth innerRef={handleRef} />);
+      assert.strictEqual(handleRef.callCount, 1);
+    });
+  });
+
   describe('browser', () => {
     it('should provide the right width to the child element', () => {
       const wrapper = mount(<EmptyWithWidth />);