Clean up package structure

src/mdx2.test.ts → src/index.test.ts

@@ -1,6 +1,6 @@
 import { dedent } from 'ts-dedent';
 import prettier from 'prettier';
-import { compileSync, SEPARATOR, wrapperJs } from './mdx2';
+import { compileSync, compile, SEPARATOR, wrapperJs } from './index';
 
 // @ts-ignore
 expect.addSnapshotSerializer({
@@ -45,7 +45,19 @@ describe('mdx2', () => {
     `);
   });
 
-  it('full snapshot', () => {
+  it('standalone jsx expressions', () => {
+    expect(
+      clean(dedent`
+        # Standalone JSX expressions
+
+        {3 + 3}
+      `)
+    ).toMatchInlineSnapshot(`const componentMeta = { includeStories: [] };`);
+  });
+});
+
+describe('full snapshots', () => {
+  it('compileSync', () => {
     const input = dedent`
       # hello
 
@@ -106,15 +118,67 @@ describe('mdx2', () => {
       export default componentMeta;
     `);
   });
+  it('compile', async () => {
+    const input = dedent`
+      # hello
 
-  it('standalone jsx expressions', () => {
-    expect(
-      clean(dedent`
-        # Standalone JSX expressions
+      <Meta title="foobar" />
 
-        {3 + 3}
-      `)
-    ).toMatchInlineSnapshot(`const componentMeta = { includeStories: [] };`);
+      world {2 + 1}
+
+      <Story name="foo">bar</Story>
+    `;
+    // @ts-ignore
+    expect(await compile(input)).toMatchInlineSnapshot(`
+/*@jsxRuntime automatic @jsxImportSource react*/
+import {Fragment as _Fragment, jsx as _jsx, jsxs as _jsxs} from "react/jsx-runtime";
+import {useMDXComponents as _provideComponents} from "@mdx-js/react";
+function MDXContent(props = {}) {
+  const {wrapper: MDXLayout} = Object.assign({}, _provideComponents(), props.components);
+  return MDXLayout ? _jsx(MDXLayout, Object.assign({}, props, {
+    children: _jsx(_createMdxContent, {})
+  })) : _createMdxContent();
+  function _createMdxContent() {
+    const _components = Object.assign({
+      h1: "h1",
+      p: "p"
+    }, _provideComponents(), props.components), {Meta, Story} = _components;
+    if (!Meta) _missingMdxReference("Meta", true);
+    if (!Story) _missingMdxReference("Story", true);
+    return _jsxs(_Fragment, {
+      children: [_jsx(_components.h1, {
+        children: "hello"
+      }), "\\n", _jsx(Meta, {
+        title: "foobar"
+      }), "\\n", _jsxs(_components.p, {
+        children: ["world ", 2 + 1]
+      }), "\\n", _jsx(Story, {
+        name: "foo",
+        children: "bar"
+      })]
+    });
+  }
+}
+function _missingMdxReference(id, component) {
+  throw new Error("Expected " + (component ? "component" : "object") + " \`" + id + "\` to be defined: you likely forgot to import, pass, or provide it.");
+}
+// =========
+export const foo = () => (
+          "bar"
+        );
+foo.storyName = 'foo';
+foo.parameters = { storySource: { source: '\\"bar\\"' } };
+
+const componentMeta = { title: 'foobar', tags: ['mdx'], includeStories: ["foo"],  };
+
+componentMeta.parameters = componentMeta.parameters || {};
+componentMeta.parameters.docs = {
+  ...(componentMeta.parameters.docs || {}),
+  page: MDXContent,
+};
+
+export default componentMeta;
+`);
   });
 });
 

src/index.ts

@@ -1,17 +1,145 @@
-import { compile as mdxCompile } from '@mdx-js/mdx';
+import { compile as mdxCompile, compileSync } from '@mdx-js/mdx';
+import generate from '@babel/generator';
+import * as t from '@babel/types';
+import cloneDeep from 'lodash/cloneDeep';
+import toBabel from 'estree-to-babel';
 import { toEstree } from 'hast-util-to-estree';
-import { plugin, postprocess } from './mdx2';
 
-export * from './mdx2';
+// Keeping as much code as possible from the original compiler to avoid breaking changes
+import {
+  genCanvasExports,
+  genStoryExport,
+  genMeta,
+  CompilerOptions,
+  Context,
+  MetaExport,
+  wrapperJs,
+  stringifyMeta,
+} from './sb-mdx-plugin';
 
-export const compile = async (code: string, { skipCsf }: { skipCsf?: boolean }) => {
-  // This async import is needed because these libraries are ESM
-  // and this file is CJS. Furthermore, we keep this file out of
-  // the src directory so that babel doesn't turn these into `require`
-  // statements when it transpiles...
+export const SEPARATOR = '// =========';
 
+export { wrapperJs };
+
+function extractExports(root: t.File, options: CompilerOptions) {
+  const context: Context = {
+    counter: 0,
+    storyNameToKey: {},
+    namedExports: {},
+  };
+  const storyExports = [];
+  const includeStories = [];
+  let metaExport: MetaExport | null = null;
+  const { code } = generate(root, {});
+  let contents: t.ExpressionStatement;
+  root.program.body.forEach((child) => {
+    if (t.isExpressionStatement(child) && t.isJSXFragment(child.expression)) {
+      if (contents) throw new Error('duplicate contents');
+      contents = child;
+    } else if (
+      t.isExportNamedDeclaration(child) &&
+      t.isVariableDeclaration(child.declaration) &&
+      child.declaration.declarations.length === 1
+    ) {
+      const declaration = child.declaration.declarations[0];
+      if (t.isVariableDeclarator(declaration) && t.isIdentifier(declaration.id)) {
+        const { name } = declaration.id;
+        context.namedExports[name] = declaration.init;
+      }
+    }
+  });
+  if (contents) {
+    const jsx = contents.expression as t.JSXFragment;
+    jsx.children.forEach((child) => {
+      if (t.isJSXElement(child)) {
+        if (t.isJSXIdentifier(child.openingElement.name)) {
+          const name = child.openingElement.name.name;
+          let stories;
+          if (['Canvas', 'Preview'].includes(name)) {
+            stories = genCanvasExports(child, context);
+          } else if (name === 'Story') {
+            stories = genStoryExport(child, context);
+          } else if (name === 'Meta') {
+            const meta = genMeta(child, options);
+            if (meta) {
+              if (metaExport) {
+                throw new Error('Meta can only be declared once');
+              }
+              metaExport = meta;
+            }
+          }
+          if (stories) {
+            Object.entries(stories).forEach(([key, story]) => {
+              includeStories.push(key);
+              storyExports.push(story);
+            });
+          }
+        }
+      } else if (t.isJSXExpressionContainer(child)) {
+        // Skip string literals & other JSX expressions
+      } else {
+        throw new Error(`Unexpected JSX child: ${child.type}`);
+      }
+    });
+  }
+
+  if (metaExport) {
+    if (!storyExports.length) {
+      storyExports.push('export const __page = () => { throw new Error("Docs-only story"); };');
+      storyExports.push('__page.parameters = { docsOnly: true };');
+      includeStories.push('__page');
+    }
+  } else {
+    metaExport = {};
+  }
+  metaExport.includeStories = JSON.stringify(includeStories);
+
+  const fullJsx = [
+    ...storyExports,
+    `const componentMeta = ${stringifyMeta(metaExport)};`,
+    wrapperJs,
+    'export default componentMeta;',
+  ].join('\n\n');
+
+  return fullJsx;
+}
+
+export const plugin = (store: any) => (root: any) => {
+  const estree = store.toEstree(root);
+  // toBabel mutates root, so we need to clone it
+  const clone = cloneDeep(estree);
+  const babel = toBabel(clone);
+  store.exports = extractExports(babel, {});
+
+  return root;
+};
+
+export const postprocess = (code: string, extractedExports: string) => {
+  const lines = code.toString().trim().split('\n');
+
+  // /*@jsxRuntime automatic @jsxImportSource react*/
+  const first = lines.shift();
+
+  return [
+    first,
+    ...lines.filter((line) => !line.match(/^export default/)),
+    SEPARATOR,
+    extractedExports,
+  ].join('\n');
+};
+
+export const mdxSync = (code: string) => {
   const store = { exports: '', toEstree };
+  const output = compileSync(code, {
+    rehypePlugins: [[plugin, store]],
+  });
+  return postprocess(output.toString(), store.exports);
+};
 
+export { mdxSync as compileSync };
+
+export const compile = async (code: string, { skipCsf }: { skipCsf?: boolean } = {}) => {
+  const store = { exports: '', toEstree };
   const output = await mdxCompile(code, {
     rehypePlugins: skipCsf ? [] : [[plugin, store]],
     providerImportSource: '@mdx-js/react',