feat(context): refactor passing navigation context (#3668)

* feat(context): refactor passing navigation context

* remove commented code in example

* adjust src/views/withNavigationFocus.js

* refactor stuff

* extract scene to variable

* Add test

* Apply CR comments

* remove junk

* bring back screen mode header

src/navigators/__tests__/StackNavigator-test.js

@@ -3,6 +3,7 @@ import { StyleSheet, View } from 'react-native';
 import renderer from 'react-test-renderer';
 
 import StackNavigator from '../createStackNavigator';
+import withNavigation from '../../views/withNavigation';
 
 const styles = StyleSheet.create({
   header: {
@@ -51,4 +52,37 @@ describe('StackNavigator', () => {
 
     expect(rendered).toMatchSnapshot();
   });
+
+  it('passes navigation to headerRight when wrapped in withNavigation', () => {
+    const spy = jest.fn();
+
+    class TestComponent extends React.Component {
+      render() {
+        return <View>{this.props.onPress(this.props.navigation)}</View>;
+      }
+    }
+
+    const TestComponentWithNavigation = withNavigation(TestComponent);
+
+    class A extends React.Component {
+      static navigationOptions = {
+        headerRight: <TestComponentWithNavigation onPress={spy} />,
+      };
+
+      render() {
+        return <View />;
+      }
+    }
+
+    const Nav = StackNavigator({ A: { screen: A } });
+
+    renderer.create(<Nav />);
+
+    expect(spy).toBeCalledWith(
+      expect.objectContaining({
+        navigate: expect.any(Function),
+        addListener: expect.any(Function),
+      })
+    );
+  });
 });

src/views/NavigationContext.js

@@ -0,0 +1,8 @@
+import React from 'react';
+import propTypes from 'prop-types';
+import createReactContext from 'create-react-context';
+
+const NavigationContext = createReactContext();
+
+export const NavigationProvider = NavigationContext.Provider;
+export const NavigationConsumer = NavigationContext.Consumer;

src/views/SceneView.js

@@ -1,19 +1,14 @@
 import React from 'react';
 import propTypes from 'prop-types';
+import { NavigationProvider } from './NavigationContext';
 
 export default class SceneView extends React.PureComponent {
-  static childContextTypes = {
-    navigation: propTypes.object.isRequired,
-  };
-
-  getChildContext() {
-    return {
-      navigation: this.props.navigation,
-    };
-  }
-
   render() {
-    const { screenProps, navigation, component: Component } = this.props;
-    return <Component screenProps={screenProps} navigation={navigation} />;
+    const { screenProps, component: Component, navigation } = this.props;
+    return (
+      <NavigationProvider value={navigation}>
+        <Component screenProps={screenProps} navigation={navigation} />
+      </NavigationProvider>
+    );
   }
 }

src/views/StackView/StackViewLayout.js

@@ -15,6 +15,7 @@ import Card from './StackViewCard';
 import Header from '../Header/Header';
 import NavigationActions from '../../NavigationActions';
 import SceneView from '../SceneView';
+import { NavigationProvider } from '../NavigationContext';
 
 import TransitionConfigs from './StackViewTransitionConfigs';
 import * as ReactNativeFeatures from '../../utils/ReactNativeFeatures';
@@ -194,9 +195,11 @@ class StackViewLayout extends React.Component {
     let floatingHeader = null;
     const headerMode = this._getHeaderMode();
     if (headerMode === 'float') {
-      floatingHeader = this._renderHeader(
-        this.props.transitionProps.scene,
-        headerMode
+      const { scene } = this.props.transitionProps;
+      floatingHeader = (
+        <NavigationProvider value={scene.descriptor.navigation}>
+          {this._renderHeader(scene, headerMode)}
+        </NavigationProvider>
       );
     }
     const {

src/views/withNavigation.js

@@ -1,24 +1,31 @@
 import React from 'react';
 import propTypes from 'prop-types';
 import hoistStatics from 'hoist-non-react-statics';
+import invariant from '../utils/invariant';
+import { NavigationConsumer } from './NavigationContext';
 
 export default function withNavigation(Component) {
   class ComponentWithNavigation extends React.Component {
     static displayName = `withNavigation(${Component.displayName ||
       Component.name})`;
 
-    static contextTypes = {
-      navigation: propTypes.object.isRequired,
-    };
-
     render() {
-      const { navigation } = this.context;
       return (
-        <Component
-          {...this.props}
-          navigation={navigation}
-          ref={this.props.onRef}
-        />
+        <NavigationConsumer>
+          {navigation => {
+            invariant(
+              !!navigation,
+              'withNavigationFocus can only be used on a view hierarchy of a navigator. The wrapped component is unable to get access to navigation from props or context.'
+            );
+            return (
+              <Component
+                {...this.props}
+                navigation={navigation}
+                ref={this.props.onRef}
+              />
+            );
+          }}
+        </NavigationConsumer>
       );
     }
   }

src/views/withNavigationFocus.js

@@ -2,26 +2,28 @@ import React from 'react';
 import propTypes from 'prop-types';
 import hoistStatics from 'hoist-non-react-statics';
 import invariant from '../utils/invariant';
+import withNavigation from './withNavigation';
 
 export default function withNavigationFocus(Component) {
   class ComponentWithNavigationFocus extends React.Component {
     static displayName = `withNavigationFocus(${Component.displayName ||
       Component.name})`;
 
-    static contextTypes = {
-      navigation: propTypes.object.isRequired,
-    };
-
-    constructor(props, context) {
-      super();
+    constructor(props) {
+      super(props);
 
       this.state = {
-        isFocused: this.getNavigation(props, context).isFocused(),
+        isFocused: props.navigation ? props.navigation.isFocused() : false,
       };
     }
 
     componentDidMount() {
-      const navigation = this.getNavigation();
+      const { navigation } = this.props;
+      invariant(
+        !!navigation,
+        'withNavigationFocus can only be used on a view hierarchy of a navigator. The wrapped component is unable to get access to navigation from props or context.'
+      );
+
       this.subscriptions = [
         navigation.addListener('didFocus', () =>
           this.setState({ isFocused: true })
@@ -36,15 +38,6 @@ export default function withNavigationFocus(Component) {
       this.subscriptions.forEach(sub => sub.remove());
     }
 
-    getNavigation = (props = this.props, context = this.context) => {
-      const navigation = props.navigation || context.navigation;
-      invariant(
-        !!navigation,
-        'withNavigationFocus can only be used on a view hierarchy of a navigator. The wrapped component is unable to get access to navigation from props or context.'
-      );
-      return navigation;
-    };
-
     render() {
       return (
         <Component
@@ -56,5 +49,5 @@ export default function withNavigationFocus(Component) {
     }
   }
 
-  return hoistStatics(ComponentWithNavigationFocus, Component);
+  return hoistStatics(withNavigation(ComponentWithNavigationFocus), Component);
 }