Some organization and style around withRouter ref.

remix-run · Aug 19, 2016 · e531f45 · e531f45
1 parent af276b2
commit e531f45
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 9 deletions.
diff --git a/modules/__tests__/withRouter-test.js b/modules/__tests__/withRouter-test.js
@@ -64,7 +64,7 @@ describe('withRouter', function () {
   })
 
   it('should support withRefs as a parameter', function (done) {
-    const WrappedApp = withRouter(App, { withRef:true })
+    const WrappedApp = withRouter(App, { withRef: true })
     const router = {
       push() {},
       replace() {},

diff --git a/modules/withRouter.js b/modules/withRouter.js
@@ -9,20 +9,23 @@ function getDisplayName(WrappedComponent) {
 }
 
 export default function withRouter(WrappedComponent, options) {
+  const { withRef } = options || {}
+
   const WithRouter = React.createClass({
     contextTypes: { router: routerShape },
     propTypes: { router: routerShape },
+
     getWrappedInstance() {
-      warning(options && options.withRef, 'To access the wrappedInstance you must provide {withRef : true} as the second argument of the withRouter call')
-      return this._wrappedComponent
+      warning(withRef, 'To access the wrappedInstance you must provide { withRef: true } as the second argument of the withRouter call')
+      return this.wrappedComponent
     },
+
     render() {
-      const router = this.props.router || this.context.router
-      if (options && options.withRef) {
-        return <WrappedComponent {...this.props} ref={(component)=>this._wrappedComponent = component} router={router} />
-      } else {
-        return <WrappedComponent {...this.props} router={router} />
-      }
+      const { router, ...props } = this.props
+
+      if (withRef) props.ref = component =>this.wrappedComponent = component
+
+      return <WrappedComponent {...props} router={router || this.context.router} />
     }
   })