Support new Context API

.gitignore

@@ -1,5 +1,7 @@
 node_modules
 .idea
 index.js
+server.js
 !src/index.js
+!src/server.js
 coverage

package.json

@@ -5,7 +5,7 @@
   "main": "index.js",
   "repository": "https://github.com/zekchan/react-ssr-error-boundary",
   "scripts": {
-    "build": "babel src/index.js > index.js",
+    "build": "babel src --out-dir . --ignore '*.test.js'",
     "test": "jest",
     "prepublish": "npm run build"
   },

src/client.test.js

@@ -94,4 +94,116 @@ describe('Client side', () => {
 
     expect(component.html()).toBe('<div><div>No errors! Context variable</div></div>')
   })
-})
+
+  it('Renders child component with new context dependencies', () => {
+    const Context = React.createContext()
+
+    function GoodComponent () {
+      return (
+        <Context.Consumer>
+          {context => <div>No errors! {context.someContext}</div>}
+        </Context.Consumer>
+      )
+    }
+
+    function ContextProvider (props) {
+      return (
+        <Context.Provider value={{someContext: 'Context variable'}}>
+          {props.children}
+        </Context.Provider>
+      )
+    }
+
+    const component = mount(
+      <ContextProvider>
+        <ErrorFallback fallBack={() => <FallBack />}>
+          <GoodComponent />
+        </ErrorFallback>
+      </ContextProvider>
+    )
+
+    expect(component.html()).toBe('<div><div>No errors! Context variable</div></div>')
+  })
+
+  it('Renders child component with multiple new context dependencies', () => {
+    const Context1 = React.createContext()
+    const Context2 = React.createContext()
+    const Context3 = React.createContext()
+
+    function GoodComponent () {
+      return (
+        <Context1.Consumer>
+          {context1 => (
+              <Context2.Consumer>
+                {context2 => (
+                    <Context3.Consumer>
+                      {context3 => (
+                          <div>No errors! {context1.someContext} {context2.someContext} {context3.someContext}</div>
+                      )}
+                    </Context3.Consumer>
+                )}
+              </Context2.Consumer>
+          )}
+        </Context1.Consumer>
+      )
+    }
+
+    function ContextProvider (props) {
+      return (
+        <Context1.Provider value={{someContext: 'Context variable1'}}>
+          <Context2.Provider value={{someContext: 'Context variable2'}}>
+            <Context3.Provider value={{someContext: 'Context variable3'}}>
+              {props.children}
+            </Context3.Provider>
+          </Context2.Provider>
+        </Context1.Provider>
+      )
+    }
+
+    const component = mount(
+      <ContextProvider>
+        <ErrorFallback fallBack={() => <FallBack />}>
+          <GoodComponent />
+        </ErrorFallback>
+      </ContextProvider>
+    )
+
+    expect(component.html()).toBe('<div><div>No errors! Context variable1 Context variable2 Context variable3</div></div>')
+  })
+
+  it('Renders fallBack component if children rendering throws error with new contexts', () => {
+    const Context = React.createContext()
+
+    function BadComponent () {
+      return (
+        <Context.Consumer>
+          {() => <BadComponentInner/>}
+        </Context.Consumer>
+      )
+    }
+
+    function BadComponentInner () {
+      throw new Error()
+    }
+
+    function ContextProvider (props) {
+      return (
+        <Context.Provider value={{someContext: 'Context variable'}}>
+          {props.children}
+        </Context.Provider>
+      )
+    }
+
+    turnOffErrors()
+    const component = mount(
+      <ContextProvider>
+        <ErrorFallback fallBack={() => <FallBack />}>
+          <BadComponent />
+        </ErrorFallback>
+      </ContextProvider>
+    )
+
+    expect(component.html()).toBe('<div><div>FallBack!</div></div>')
+    turnOnErrors()
+  })
+})

src/index.js

@@ -1,19 +1,9 @@
 import React, { Component } from 'react'
 
-const isServer = typeof window === 'undefined'
+const server = typeof window === 'undefined' && require('./server')
 
 export function withContext (contextTypes = {}) {
-  class ProvideContext extends Component {
-    static childContextTypes = contextTypes
-
-    getChildContext () {
-      return this.props.context
-    }
-
-    render () {
-      return this.props.children
-    }
-  }
+  const ProvideContext = server && server._makeProvider(contextTypes)
 
   class ServerBoundary extends Component {
     static defaultProps = {
@@ -31,19 +21,7 @@ export function withContext (contextTypes = {}) {
     }
 
     render () {
-      if (isServer) {
-        try {
-          const __html = require('react-dom/server').renderToStaticMarkup(
-            <ProvideContext context={this.context}>
-              {this.props.children}
-            </ProvideContext>
-          )
-          return <div dangerouslySetInnerHTML={{__html}} />
-        } catch (e) {
-          return <div>{this.props.fallBack()}</div>
-        }
-      }
-
+      if (server) return server._render(this, ProvideContext)
       return <div>{this.state.elementToRender}</div>
     }
   }

src/server.js

@@ -0,0 +1,108 @@
+import React, { Component } from 'react'
+import { renderToStaticMarkup } from 'react-dom/server'
+
+// Array to store React contexts
+const contexts = []
+
+// Use specific context
+export function useContext (Context) {
+  contexts.push(Context);
+}
+
+// Clear contexts in use
+export function clearContexts () {
+  contexts.length = 0
+}
+
+// Shim `React.createContext()` to capture contexts
+let shimmed = false
+
+export function shim () {
+  if (shimmed) return
+  shimmed = true
+
+  const { createContext } = React
+  React.createContext = function() {
+    const Context = createContext.apply(this, arguments)
+    contexts.push(Context)
+    return Context
+  }
+}
+
+// Create Provider component for specific context types
+export function _makeProvider (contextTypes) {
+  return class ProvideContext extends Component {
+    static childContextTypes = contextTypes
+
+    getChildContext () {
+      return this.props.context
+    }
+
+    render () {
+      return this.props.children
+    }
+  }
+}
+
+// Server-side render, propagating contexts to children
+export function _render (self, ProvideContext) {
+  // Wrap element in provider for legacy context
+  const element = (
+    <ProvideContext context={self.context}>
+      {self.props.children}
+    </ProvideContext>
+  )
+
+  // Wrap element in consumers and providers for new context.
+  // Consumers are added outside HTML render, to capture value of all contexts.
+  // Providers are added inside HTML render, to reinject context values back
+  // into the rendered element.
+  return stackConsumers(values => {
+    const elementWithProviders = stackProviders(element, values)
+
+    try {
+      const __html = renderToStaticMarkup(elementWithProviders)
+      return <div dangerouslySetInnerHTML={{__html}} />
+    } catch (e) {
+      return <div>{self.props.fallBack()}</div>
+    }
+  })
+}
+
+// Create stack of consumer components
+function stackConsumers (retFn) {
+  const values = []
+  let fn = () => retFn(values)
+
+  for (let i = 0; i < contexts.length; i++) {
+    const Context = contexts[i]
+    const childFn = fn
+
+    fn = () => (
+      <Context.Consumer>
+        {value => {
+          values[i] = value
+          return childFn()
+        }}
+      </Context.Consumer>
+    )
+  }
+
+  return fn()
+}
+
+// Create stack of provider components
+function stackProviders (element, values) {
+  for (let i = 0; i < values.length; i++) {
+    const Context = contexts[i]
+    const childElement = element
+
+    element = (
+      <Context.Provider value={values[i]}>
+        {childElement}
+      </Context.Provider>
+    )
+  }
+
+  return element
+}