Adding support to get the component instance that withRouter HOC wraps

modules/__tests__/withRouter-test.js

@@ -12,6 +12,9 @@ describe('withRouter', function () {
     propTypes: {
       router: routerShape.isRequired
     }
+    testFunction() {
+      return 'hello from the test function'
+    }
     render() {
       expect(this.props.router).toExist()
       return <h1>App</h1>
@@ -28,7 +31,7 @@ describe('withRouter', function () {
   })
 
   it('puts router on context', function (done) {
-    const WrappedApp = withRouter(App)
+    const WrappedApp = withRouter()(App)
 
     render((
       <Router history={createHistory('/')}>
@@ -40,7 +43,7 @@ describe('withRouter', function () {
   })
 
   it('still uses router prop if provided', function (done) {
-    const Test = withRouter(function (props) {
+    const Test = withRouter()(function (props) {
       props.test(props)
       return null
     })
@@ -59,4 +62,20 @@ describe('withRouter', function () {
 
     render(<Test router={router} test={test} />, node, done)
   })
+
+  it('should support withRefs as a parameter',function (done) {
+    const WrappedApp = withRouter({ withRef:true })(App)
+    const router = {
+      push() {},
+      replace() {},
+      go() {},
+      goBack() {},
+      goForward() {},
+      setRouteLeaveHook() {},
+      isActive() {}
+    }
+    const component = render((<WrappedApp router={router}/>), node, done)
+    expect(component.getWrappedInstance().testFunction()).toEqual('hello from the test function')
+  })
+
 })

modules/withRouter.js

@@ -1,23 +1,40 @@
 import React from 'react'
 import hoistStatics from 'hoist-non-react-statics'
 import { routerShape } from './PropTypes'
+import warning from './routerWarning'
 
 function getDisplayName(WrappedComponent) {
   return WrappedComponent.displayName || WrappedComponent.name || 'Component'
 }
+export default function withRouter(options) {
+  const wrapWithRouter = function (WrappedComponent) {
+    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 first argument of the withRouter call')
+        return this.refs.wrappedInstance
+      },
+      render() {
+        const router = this.props.router || this.context.router
+        if(options && options.withRef) {
+          return <WrappedComponent {...this.props} ref="wrappedInstance" router={router} />
+        }else{
+          return <WrappedComponent {...this.props} router={router} />
+        }
+      }
+    })
+    WithRouter.displayName = `withRouter(${getDisplayName(WrappedComponent)})`
+    WithRouter.WrappedComponent = WrappedComponent
 
-export default function withRouter(WrappedComponent) {
-  const WithRouter = React.createClass({
-    contextTypes: { router: routerShape },
-    propTypes: { router: routerShape },
-    render() {
-      const router = this.props.router || this.context.router
-      return <WrappedComponent {...this.props} router={router} />
-    }
-  })
+    return hoistStatics(WithRouter, WrappedComponent)
+  }
 
-  WithRouter.displayName = `withRouter(${getDisplayName(WrappedComponent)})`
-  WithRouter.WrappedComponent = WrappedComponent
-
-  return hoistStatics(WithRouter, WrappedComponent)
+  if(typeof(options) === 'function') {
+    warning(false,'passing a component to the first invocation of withRouter has been depreciated, withRouter'+ 
+    'now needs to be invoked twice, once to pass the options through, and a second time with the component eg. withRouter()(MyComponent) or withRouter({withRef:true})(MyComponent)')
+    return wrapWithRouter(options)
+  }else{
+    return wrapWithRouter
+  }
 }