From 925028c3af6d53b84cbe9e932f78de4bd8a1087b Mon Sep 17 00:00:00 2001
From: itelofh <itelofilho@gmail.com>
Date: Sun, 22 Jul 2018 12:39:16 -0300
Subject: [PATCH] add innerRef to withWidth and withTheme

---
 packages/material-ui/src/styles/withTheme.d.ts |  1 +
 packages/material-ui/src/styles/withTheme.js   | 11 ++++++++++-
 .../material-ui/src/styles/withTheme.test.js   | 16 ++++++++++++++++
 .../material-ui/src/withWidth/withWidth.d.ts   |  1 +
 .../material-ui/src/withWidth/withWidth.js     |  8 ++++++--
 .../src/withWidth/withWidth.test.js            | 18 +++++++++++++++++-
 6 files changed, 51 insertions(+), 4 deletions(-)

diff --git a/packages/material-ui/src/styles/withTheme.d.ts b/packages/material-ui/src/styles/withTheme.d.ts
index 7970bfe2610c5a..71d43b1c8a67ab 100644
--- a/packages/material-ui/src/styles/withTheme.d.ts
+++ b/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>>(
diff --git a/packages/material-ui/src/styles/withTheme.js b/packages/material-ui/src/styles/withTheme.js
index bedd447f25fd11..0c18daae0f50b2 100644
--- a/packages/material-ui/src/styles/withTheme.js
+++ b/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') {
diff --git a/packages/material-ui/src/styles/withTheme.test.js b/packages/material-ui/src/styles/withTheme.test.js
index 743034a65f2bb0..ff1ba991ed9fc8 100644
--- a/packages/material-ui/src/styles/withTheme.test.js
+++ b/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);
+    });
+  });
 });
diff --git a/packages/material-ui/src/withWidth/withWidth.d.ts b/packages/material-ui/src/withWidth/withWidth.d.ts
index fa980f3deef647..5ded5189a33bf8 100644
--- a/packages/material-ui/src/withWidth/withWidth.d.ts
+++ b/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(
diff --git a/packages/material-ui/src/withWidth/withWidth.js b/packages/material-ui/src/withWidth/withWidth.js
index 670c3d36c2dc10..e2efb9d82b573a 100644
--- a/packages/material-ui/src/withWidth/withWidth.js
+++ b/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
      */
diff --git a/packages/material-ui/src/withWidth/withWidth.test.js b/packages/material-ui/src/withWidth/withWidth.test.js
index 8902ebc344fd5f..f63310ee8b5b02 100644
--- a/packages/material-ui/src/withWidth/withWidth.test.js
+++ b/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 />);