Add useMDXComponents hook

packages/mdx/create-element.js

@@ -1,5 +1,5 @@
 const React = require('react')
-const {withMDXComponents} = require('@mdx-js/tag')
+const {useMDXComponents} = require('@mdx-js/tag')
 
 const TYPE_PROP_NAME = '__MDX_TYPE_PLEASE_DO_NOT_USE__'
 
@@ -8,12 +8,13 @@ const DEFAULTS = {
   wrapper: 'div'
 }
 
-const MDXCreateElementInner = ({
-  components = {},
+const MDXCreateElement = ({
+  components: propComponents,
   __MDX_TYPE_PLEASE_DO_NOT_USE__,
   parentName,
   ...etc
 }) => {
+  const components = useMDXComponents(propComponents)
   const type = __MDX_TYPE_PLEASE_DO_NOT_USE__
   const Component =
     components[`${parentName}.${type}`] ||
@@ -23,9 +24,6 @@ const MDXCreateElementInner = ({
 
   return React.createElement(Component, etc)
 }
-MDXCreateElementInner.displayName = 'MDXCreateElementInner'
-
-const MDXCreateElement = withMDXComponents(MDXCreateElementInner)
 MDXCreateElement.displayName = 'MDXCreateElement'
 
 module.exports = function(type, props) {

packages/mdx/test/index.test.js

@@ -396,6 +396,7 @@ export default class MDXContent extends React.Component {
              <Layout {...layoutProps} {...props}>
 
 
+
 <h1 >{\`Hello, world!\`}</h1>
 <p >{\`I'm an awesome paragraph.\`}</p>
 {/* I'm a comment */}

packages/tag/src/index.js

@@ -2,5 +2,6 @@ export {default as MDXTag} from './mdx-tag'
 export {
   default as MDXContext,
   MDXProvider,
-  withMDXComponents
+  withMDXComponents,
+  useMDXComponents
 } from './mdx-context'

packages/tag/src/mdx-context.js

@@ -5,17 +5,25 @@ const isFunction = obj => typeof obj === 'function'
 const MDXContext = React.createContext({})
 
 export const withMDXComponents = Component => props => {
-  const contextComponents = React.useContext(MDXContext)
-  const allComponents = props.components || contextComponents
+  const allComponents = useMDXComponents(props.components)
 
   return <Component {...props} components={allComponents} />
 }
 
+export const useMDXComponents = components => {
+  const contextComponents = React.useContext(MDXContext)
+  let allComponents = contextComponents
+  if (components) {
+    allComponents = isFunction(components)
+      ? components(contextComponents)
+      : components
+  }
+
+  return allComponents
+}
+
 export const MDXProvider = props => {
-  const outerContextComponents = React.useContext(MDXContext)
-  const allComponents = isFunction(props.components)
-    ? props.components(outerContextComponents)
-    : props.components
+  const allComponents = useMDXComponents(props.components)
 
   return (
     <MDXContext.Provider value={allComponents}>

packages/tag/src/mdx-tag.js

@@ -1,39 +1,36 @@
-import React, {Component} from 'react'
-import {withMDXComponents} from './mdx-context'
+import React from 'react'
+import {useMDXComponents} from './mdx-context'
 
 const defaults = {
   inlineCode: 'code',
   wrapper: 'div'
 }
 
-class MDXTag extends Component {
-  render() {
-    const {
-      name,
-      parentName,
-      props: childProps = {},
-      children,
-      components = {},
-      Layout,
-      layoutProps
-    } = this.props
+const MDXTag = ({
+  name,
+  parentName,
+  props: childProps = {},
+  children,
+  components: propComponents,
+  Layout,
+  layoutProps
+}) => {
+  const components = useMDXComponents(propComponents)
+  const Component =
+    components[`${parentName}.${name}`] ||
+    components[name] ||
+    defaults[name] ||
+    name
 
-    const Component =
-      components[`${parentName}.${name}`] ||
-      components[name] ||
-      defaults[name] ||
-      name
-
-    if (Layout) {
-      return (
-        <Layout components={components} {...layoutProps}>
-          <Component {...childProps}>{children}</Component>
-        </Layout>
-      )
-    }
-
-    return <Component {...childProps}>{children}</Component>
+  if (Layout) {
+    return (
+      <Layout components={components} {...layoutProps}>
+        <Component {...childProps}>{children}</Component>
+      </Layout>
+    )
   }
+
+  return <Component {...childProps}>{children}</Component>
 }
 
-export default withMDXComponents(MDXTag)
+export default MDXTag