withNavigation on stateless components does not work. add withRef option

Closes #447

fbshipit-source-id: f696928

src/ExNavigationComponents.js

@@ -6,6 +6,7 @@ import React, { PropTypes } from 'react';
 
 import UUID from 'uuid-js';
 
+import invariant from 'invariant';
 import { createSelector } from 'reselect';
 import hoistStatics from 'hoist-non-react-statics';
 import PureComponent from './utils/PureComponent';
@@ -156,7 +157,10 @@ import {
 
 const NavigatorPropType = PropTypes.object;
 
-export function withNavigation<T>(WrappedComponent: ReactClass<T>) {
+export function withNavigation<T>(
+  WrappedComponent: ReactClass<T>,
+  { withRef }
+) {
   class WithNavigation extends PureComponent {
     _wrappedInstance: ReactElement<T>;
 
@@ -173,9 +177,7 @@ export function withNavigation<T>(WrappedComponent: ReactClass<T>) {
     render() {
       return (
         <WrappedComponent
-          ref={c => {
-            this._wrappedInstance = c;
-          }}
+          ref={withRef ? this.setWrappedInstance : undefined}
           navigation={this.getNavigationContext()}
           navigator={this.getCurrentNavigator()}
           {...this.props}
@@ -191,9 +193,20 @@ export function withNavigation<T>(WrappedComponent: ReactClass<T>) {
     }
 
     getWrappedInstance() {
+      if (__DEV__) {
+        invariant(
+          withRef,
+          'To access the wrapped instance, you need to specify ' +
+            '{ withRef: true } in the options argument of withNavigation call.'
+        );
+      }
       return this._wrappedInstance;
     }
 
+    setWrappedInstance(ref) {
+      this._wrappedInstance = ref;
+    }
+
     getNavigationContext() {
       return this.props.navigation || this.context.navigation;
     }