From 24964b42330d34e9189c75a156fbbd12ef939773 Mon Sep 17 00:00:00 2001
From: Chinmay Chaudhary <chaudharychinmay@gmail.com>
Date: Thu, 22 Jul 2021 19:55:15 +0530
Subject: [PATCH] Use forwardRef to pass on the ref (#182)

---
 package.json           |  6 +++---
 src/withStyles.js      | 21 +++++++++++++++++----
 test/insertCss.test.js |  5 ++++-
 3 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/package.json b/package.json
index 0c53feb5..87eced0a 100644
--- a/package.json
+++ b/package.json
@@ -33,10 +33,11 @@
   "main": "src/index.js",
   "dependencies": {
     "hoist-non-react-statics": "^3.0.0",
-    "loader-utils": "^1.2.3"
+    "loader-utils": "^1.2.3",
+    "prop-types": "^15.7.2"
   },
   "peerDependencies": {
-    "react": "^16.3.0",
+    "react": "^16.8.2",
     "react-dom": "^16.8.2"
   },
   "devDependencies": {
@@ -56,7 +57,6 @@
     "husky": "^1.3.1",
     "jest": "^24.1.0",
     "prettier": "^1.16.4",
-    "prop-types": "^15.7.2",
     "react": "^16.8.2",
     "react-dom": "^16.8.2",
     "rollup": "^1.2.1",
diff --git a/src/withStyles.js b/src/withStyles.js
index 5b7d7d73..5bb9b56c 100644
--- a/src/withStyles.js
+++ b/src/withStyles.js
@@ -8,6 +8,7 @@
  */
 
 import React from 'react'
+import PropTypes from 'prop-types'
 import hoistStatics from 'hoist-non-react-statics'
 
 import StyleContext from './StyleContext'
@@ -27,17 +28,29 @@ function withStyles(...styles) {
       }
 
       render() {
-        return <ComposedComponent {...this.props} />
+        const { __$$withStylesRef, ...props } = this.props
+        return <ComposedComponent ref={__$$withStylesRef} {...props} />
       }
     }
 
     const displayName = ComposedComponent.displayName || ComposedComponent.name || 'Component'
 
-    WithStyles.displayName = `WithStyles(${displayName})`
+    WithStyles.propTypes = {
+      __$$withStylesRef: PropTypes.func,
+    }
+    WithStyles.defaultProps = {
+      __$$withStylesRef: undefined,
+    }
     WithStyles.contextType = StyleContext
-    WithStyles.ComposedComponent = ComposedComponent
 
-    return hoistStatics(WithStyles, ComposedComponent)
+    const ForwardedWithStyles = React.forwardRef((props, ref) => (
+      <WithStyles {...props} __$$withStylesRef={ref} />
+    ))
+
+    ForwardedWithStyles.ComposedComponent = ComposedComponent
+    ForwardedWithStyles.displayName = `WithStyles(${displayName})`
+
+    return hoistStatics(ForwardedWithStyles, ComposedComponent)
   }
 }
 
diff --git a/test/insertCss.test.js b/test/insertCss.test.js
index 7e46e923..dc51d0eb 100644
--- a/test/insertCss.test.js
+++ b/test/insertCss.test.js
@@ -25,7 +25,10 @@ describe('insertCss(styles, options)', () => {
   it('Should insert and remove multiple <style> elements for a single module', () => {
     const css1 = 'body { color: red; }'
     const css2 = 'body { color: blue; }'
-    const removeCss = insertCss([[1, css1], [1, css2]])
+    const removeCss = insertCss([
+      [1, css1],
+      [1, css2],
+    ])
     let style = global.document.getElementsByTagName('style')
     expect(style).toHaveLength(2)
     expect(style[0].textContent).toBe(css1)