From a2707cb245a3b584fe55aa5b800fba5dde919efd Mon Sep 17 00:00:00 2001
From: Bodo Graumann <mail@bodograumann.de>
Date: Tue, 31 Oct 2023 15:44:18 +0100
Subject: [PATCH] Remove exports from mdx files

---
 .../transforms/__tests__/mdx-to-csf.test.ts   |  7 ------
 code/lib/codemod/src/transforms/mdx-to-csf.ts | 23 +++++++++++++++++++
 2 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/code/lib/codemod/src/transforms/__tests__/mdx-to-csf.test.ts b/code/lib/codemod/src/transforms/__tests__/mdx-to-csf.test.ts
index 2cd487bb892c..669a836e65ea 100644
--- a/code/lib/codemod/src/transforms/__tests__/mdx-to-csf.test.ts
+++ b/code/lib/codemod/src/transforms/__tests__/mdx-to-csf.test.ts
@@ -247,9 +247,6 @@ test('extract esm into csf head code', () => {
 
     # hello
 
-    export const args = { bla: 1 };
-    export const Template = (args) => <Button {...args} />;
-
     <Meta of={FoobarStories} />
 
     world {2 + 1}
@@ -412,8 +409,6 @@ test('duplicate story name', () => {
     import { Button } from './Button';
     import * as FoobarStories from './Foobar.stories';
 
-    export const Default = (args) => <Button {...args} />;
-
     <Meta of={FoobarStories} />
 
     <Story of={FoobarStories.Default_} />
@@ -502,8 +497,6 @@ test('kebab case file name', () => {
     import { Kebab } from './my-component/some-kebab-case';
     import * as SomeKebabCaseStories from './some-kebab-case.stories';
 
-    export const Template = (args) => <Kebab {...args} />;
-
     <Meta of={SomeKebabCaseStories} />
 
     <Story of={SomeKebabCaseStories.MuchKebab} />
diff --git a/code/lib/codemod/src/transforms/mdx-to-csf.ts b/code/lib/codemod/src/transforms/mdx-to-csf.ts
index f05874dfb69a..06cb563bd333 100644
--- a/code/lib/codemod/src/transforms/mdx-to-csf.ts
+++ b/code/lib/codemod/src/transforms/mdx-to-csf.ts
@@ -4,6 +4,7 @@ import { babelParse, babelParseExpression } from '@storybook/csf-tools';
 import { remark } from 'remark';
 import type { Root } from 'remark-mdx';
 import remarkMdx from 'remark-mdx';
+import type { Parent } from 'unist';
 import { SKIP, visit } from 'unist-util-visit';
 import { is } from 'unist-util-is';
 import type {
@@ -205,6 +206,7 @@ export function transform(source: string, baseName: string): [mdx: string, csf:
     return [mdxProcessor.stringify(root), null];
   }
 
+  cleanUpMdx(root);
   addStoriesImport(root, baseName, storyNamespaceName);
 
   const newStatements: t.Statement[] = [
@@ -335,6 +337,27 @@ function getEsmAst(root: Root) {
   return file;
 }
 
+function cleanUpMdx(root: Root): void {
+  visit(root, ['mdxjsEsm'], (node: MdxjsEsm, index?: number, parent?: Parent) => {
+    // @ts-expect-error File is not yet exposed, see https://github.com/babel/babel/issues/11350#issuecomment-644118606
+    const file: BabelFile = new babel.File(
+      { filename: 'info.path' },
+      { code: node.value, ast: babelParse(node.value) }
+    );
+
+    file.path.traverse({
+      ExportDeclaration: (path) => path.remove(),
+    });
+    const code = recast.print(file.path.node).code.trim();
+    if (code === '') {
+      parent.children.splice(index, 1);
+      return index;
+    }
+    node.value = code;
+    return index + 1;
+  });
+}
+
 function addStoriesImport(root: Root, baseName: string, storyNamespaceName: string): void {
   let found = false;