refactor: rework theme provider and HOC

Author: satya164

.flowconfig

@@ -35,6 +35,10 @@
 .*/docs/node_modules/fbjs/.*
 .*/docs/node_modules/react/.*
 
+; Ignore some modules we don't need to parse
+.*/node_modules/prettier/.*
+.*/node_modules/eslint.*
+
 [include]
 
 [libs]
@@ -66,4 +70,4 @@ suppress_comment=\\(.\\|\n\\)*\\$FlowExpectedError
 unsafe.enable_getters_and_setters=true
 
 [version]
-^0.40.0
+^0.40.0

src/core/ThemeProvider.js

@@ -14,7 +14,7 @@ type Props = {
   theme?: Theme,
 };
 
-export const theme = 'react-native-paper$theme';
+export const channel = 'react-native-paper$theme';
 
 export default class ThemeProvider
   extends PureComponent<DefaultProps, Props, void> {
@@ -28,15 +28,41 @@ export default class ThemeProvider
   };
 
   static childContextTypes = {
-    [theme]: PropTypes.object,
+    [channel]: PropTypes.object,
   };
 
   getChildContext() {
     return {
-      [theme]: this.props.theme,
+      [channel]: {
+        subscribe: this._subscribe,
+        get: this._get,
+      },
     };
   }
 
+  componentWillReceiveProps(nextProps: *) {
+    if (this.props.theme !== nextProps.theme) {
+      this._subscriptions.forEach(cb => cb(nextProps.theme));
+    }
+  }
+
+  _subscriptions = [];
+
+  _subscribe = (callback: Theme => void) => {
+    this._subscriptions.push(callback);
+
+    const remove = () => {
+      const index = this._subscriptions.indexOf(callback);
+      if (index > -1) {
+        this._subscriptions.splice(index, 1);
+      }
+    };
+
+    return { remove };
+  };
+
+  _get = () => this.props.theme;
+
   render() {
     return Children.only(this.props.children);
   }

src/core/withTheme.js

@@ -2,14 +2,16 @@
 
 import React, { PureComponent } from 'react';
 import PropTypes from 'prop-types';
-import { theme } from './ThemeProvider';
+import { channel } from './ThemeProvider';
 import type { Theme } from '../types/Theme';
 
 type State = {
   theme: Theme,
 };
 
-export default function withTheme<T: *>(Comp: ReactClass<T>): ReactClass<T> {
+export default function withThemeName<T: *>(
+  Comp: ReactClass<T>
+): ReactClass<T> {
   class ThemedComponent extends PureComponent<void, *, State> {
     static displayName = `withTheme(${Comp.displayName || Comp.name})`;
 
@@ -18,38 +20,55 @@ export default function withTheme<T: *>(Comp: ReactClass<T>): ReactClass<T> {
     };
 
     static contextTypes = {
-      [theme]: PropTypes.object,
+      [channel]: PropTypes.object,
     };
 
     constructor(props, context) {
       super(props, context);
+
+      const theme = this.context[channel] && this.context[channel].get();
+
+      if (typeof theme !== 'object' && typeof this.props.theme !== 'object') {
+        throw new Error(
+          `Couldn't find theme in the context or props. ` +
+            `You need to wrap your component in '<ThemeProvider />' or pass a 'theme' prop`
+        );
+      }
+
       this.state = {
-        theme: this._merge(context[theme], props.theme),
+        theme: this._merge(theme, this.props.theme),
       };
     }
 
     state: State;
 
     componentDidMount() {
-      if (typeof this.state.theme !== 'object' || this.state.theme === null) {
-        throw new Error(
-          "Couldn't find theme in the context or props. " +
-            "You need to wrap your component in '<ThemeProvider />' or pass a 'theme' prop"
+      this._subscription =
+        this.context[channel] &&
+        this.context[channel].subscribe(theme =>
+          this.setState({ theme: this._merge(theme, this.props.theme) })
         );
-      }
     }
 
-    componentWillReceiveProps(nextProps, nextContext: any) {
-      if (
-        nextProps.theme !== this.props.theme ||
-        nextContext[theme] !== this.context[theme]
-      ) {
+    componentWillReceiveProps(nextProps: *) {
+      if (this.props.theme !== nextProps.theme) {
         this.setState({
-          theme: this._merge(nextContext[theme], nextProps.theme),
+          theme: this._merge(
+            this.context[channel] && this.context[channel].get(),
+            nextProps.theme
+          ),
         });
       }
     }
 
+    componentWillUnmount() {
+      this._subscription && this._subscription.remove();
+    }
+
+    getWrappedInstance() {
+      return this._root;
+    }
+
     setNativeProps(...args) {
       return this._root.setNativeProps(...args);
     }
@@ -62,6 +81,7 @@ export default function withTheme<T: *>(Comp: ReactClass<T>): ReactClass<T> {
       }
     };
 
+    _subscription: Object;
     _root: any;
 
     render() {