Codemods: Enhance mdx-to-csf codemod

code/lib/codemod/src/index.ts

@@ -87,17 +87,17 @@ export async function runCodemod(codemod: any, { glob, logger, dryRun, rename, p
         shell: true,
       }
     );
-    if (result.status === 1) {
+
+    if (codemod === 'mdx-to-csf' && result.status === 1) {
+      logger.log(
+        'The codemod was not able to transform the files mentioned above. We have renamed the files to .mdx.broken. Please check the files and rename them back to .mdx after you have either manually transformed them to mdx + csf or fixed the issues so that the codemod can transform them.'
+      );
+    } else if (result.status === 1) {
       logger.log('Skipped renaming because of errors.');
       return;
     }
   }
 
-  if (!renameParts && codemod === 'mdx-to-csf') {
-    renameParts = ['.stories.mdx', '.mdx'];
-    rename = '.stories.mdx:.mdx;';
-  }
-
   if (renameParts) {
     const [from, to] = renameParts;
     logger.log(`=> Renaming ${rename}: ${files.length} files`);

code/lib/codemod/src/transforms/__tests__/mdx-to-csf.test.ts

@@ -598,6 +598,50 @@ it('story child is identifier', async () => {
   `);
 });
 
+it('should replace ArgsTable by Controls', async () => {
+  const input = dedent`
+      import { ArgsTable } from '@storybook/addon-docs/blocks';
+      import { Button } from './button';
+
+      Dummy Code 
+
+      <ArgsTable of="string" />
+    `;
+
+  const mdx = await jscodeshift({ source: input, path: 'Foobar.stories.mdx' });
+
+  expect(mdx).toMatchInlineSnapshot(`
+    import { Controls } from '@storybook/blocks';
+    import { Button } from './button';
+
+    Dummy Code
+
+    <Controls />
+  `);
+});
+
+it('should not create stories.js file if there are no components', async () => {
+  const input = dedent`
+  import { Meta } from '@storybook/addon-docs';
+
+  <Meta title='Example/Introduction' />
+
+  # Welcome to Storybook
+    `;
+
+  const mdx = await jscodeshift({ source: input, path: 'Foobar.stories.mdx' });
+
+  expect(fs.writeFileSync).not.toHaveBeenCalled();
+
+  expect(mdx).toMatchInlineSnapshot(`
+  import { Meta } from '@storybook/blocks';
+
+  <Meta title="Example/Introduction" />
+
+  # Welcome to Storybook
+  `);
+});
+
 it('nameToValidExport', () => {
   expect(nameToValidExport('1 starts with digit')).toMatchInlineSnapshot(`$1StartsWithDigit`);
   expect(nameToValidExport('name')).toMatchInlineSnapshot(`Name`);

code/lib/codemod/src/transforms/mdx-to-csf.ts

@@ -24,6 +24,9 @@ import type { MdxFlowExpression } from 'mdast-util-mdx-expression';
 
 const mdxProcessor = remark().use(remarkMdx) as ReturnType<typeof remark>;
 
+const renameList: { original: string; baseName: string }[] = [];
+const brokenList: { original: string; baseName: string }[] = [];
+
 export default async function jscodeshift(info: FileInfo) {
   const parsed = path.parse(info.path);
 
@@ -37,18 +40,37 @@ export default async function jscodeshift(info: FileInfo) {
     baseName += '_';
   }
 
-  const result = await transform(info, path.basename(baseName));
+  try {
+    const { csf, mdx } = await transform(info, path.basename(baseName));
+
+    if (csf != null) {
+      fs.writeFileSync(`${baseName}.stories.js`, csf);
+    }
 
-  const [mdx, csf] = result;
+    renameList.push({ original: info.path, baseName });
 
-  if (csf != null) {
-    fs.writeFileSync(`${baseName}.stories.js`, csf);
+    return mdx;
+  } catch (e) {
+    brokenList.push({ original: info.path, baseName });
+    throw e;
   }
-
-  return mdx;
 }
 
-export async function transform(info: FileInfo, baseName: string): Promise<[string, string]> {
+// The JSCodeshift CLI doesn't return a list of files that were transformed or skipped.
+// This is a workaround to rename the files after the transformation, which we can remove after we switch from jscodeshift to another solution.
+process.on('exit', () => {
+  renameList.forEach((file) => {
+    fs.renameSync(file.original, `${file.baseName}.mdx`);
+  });
+  brokenList.forEach((file) => {
+    fs.renameSync(file.original, `${file.original}.broken`);
+  });
+});
+
+export async function transform(
+  info: FileInfo,
+  baseName: string
+): Promise<{ mdx: string; csf: string | null }> {
   const root = mdxProcessor.parse(info.source);
   const storyNamespaceName = nameToValidExport(`${baseName}Stories`);
 
@@ -74,25 +96,42 @@ export async function transform(info: FileInfo, baseName: string): Promise<[stri
     node.value = node.value
       .replaceAll('@storybook/addon-docs/blocks', '@storybook/blocks')
       .replaceAll('@storybook/addon-docs', '@storybook/blocks');
+
+    if (node.value.includes('@storybook/blocks')) {
+      // @ts-ignore
+      const file: BabelFile = new babel.File(
+        { filename: 'info.path' },
+        { code: node.value, ast: babelParse(node.value) }
+      );
+
+      file.path.traverse({
+        ImportDeclaration(path) {
+          if (path.node.source.value === '@storybook/blocks') {
+            path.get('specifiers').forEach((specifier) => {
+              if (specifier.isImportSpecifier()) {
+                const imported = specifier.get('imported');
+                if (imported.isIdentifier() && imported.node.name === 'ArgsTable') {
+                  imported.node.name = 'Controls';
+                }
+              }
+            });
+          }
+        },
+      });
+
+      node.value = recast.print(file.ast).code;
+    }
   });
 
   const file = getEsmAst(root);
 
   visit(root, ['mdxJsxFlowElement', 'mdxJsxTextElement'], (node, index, parent) => {
     if (node.type === 'mdxJsxFlowElement' || node.type === 'mdxJsxTextElement') {
-      if (is(node, { name: 'Meta' })) {
-        metaAttributes.push(...node.attributes);
-        node.attributes = [
-          {
-            type: 'mdxJsxAttribute',
-            name: 'of',
-            value: {
-              type: 'mdxJsxAttributeValueExpression',
-              value: storyNamespaceName,
-            },
-          },
-        ];
+      if (is(node, { name: 'ArgsTable' })) {
+        node.name = 'Controls';
+        node.attributes = [];
       }
+
       if (is(node, { name: 'Story' })) {
         const nameAttribute = node.attributes.find(
           (it) => it.type === 'mdxJsxAttribute' && it.name === 'name'
@@ -167,21 +206,6 @@ export async function transform(info: FileInfo, baseName: string): Promise<[stri
     return undefined;
   });
 
-  const metaProperties = metaAttributes.flatMap((attribute) => {
-    if (attribute.type === 'mdxJsxAttribute') {
-      if (typeof attribute.value === 'string') {
-        return [t.objectProperty(t.identifier(attribute.name), t.stringLiteral(attribute.value))];
-      }
-      return [
-        t.objectProperty(
-          t.identifier(attribute.name),
-          babelParseExpression(attribute.value?.value ?? '') as any as t.Expression
-        ),
-      ];
-    }
-    return [];
-  });
-
   file.path.traverse({
     // remove mdx imports from csf
     ImportDeclaration(path) {
@@ -196,11 +220,49 @@ export async function transform(info: FileInfo, baseName: string): Promise<[stri
     },
   });
 
-  if (storiesMap.size === 0 && metaAttributes.length === 0) {
+  if (storiesMap.size === 0) {
     // A CSF file must have at least one story, so skip migrating if this is the case.
-    return [mdxProcessor.stringify(root), ''];
+    return {
+      csf: null,
+      mdx: mdxProcessor.stringify(root),
+    };
   }
 
+  // Rewrites the Meta tag to use the new story namespace
+  visit(root, ['mdxJsxFlowElement', 'mdxJsxTextElement'], (node, index, parent) => {
+    if (
+      (node.type === 'mdxJsxFlowElement' || node.type === 'mdxJsxTextElement') &&
+      is(node, { name: 'Meta' })
+    ) {
+      metaAttributes.push(...node.attributes);
+      node.attributes = [
+        {
+          type: 'mdxJsxAttribute',
+          name: 'of',
+          value: {
+            type: 'mdxJsxAttributeValueExpression',
+            value: storyNamespaceName,
+          },
+        },
+      ];
+    }
+  });
+
+  const metaProperties = metaAttributes.flatMap((attribute) => {
+    if (attribute.type === 'mdxJsxAttribute') {
+      if (typeof attribute.value === 'string') {
+        return [t.objectProperty(t.identifier(attribute.name), t.stringLiteral(attribute.value))];
+      }
+      return [
+        t.objectProperty(
+          t.identifier(attribute.name),
+          babelParseExpression(attribute.value?.value ?? '') as any as t.Expression
+        ),
+      ];
+    }
+    return [];
+  });
+
   addStoriesImport(root, baseName, storyNamespaceName);
 
   const newStatements: t.Statement[] = [
@@ -297,7 +359,10 @@ export async function transform(info: FileInfo, baseName: string): Promise<[stri
     filepath: path,
   });
 
-  return [newMdx, output];
+  return {
+    csf: output,
+    mdx: newMdx,
+  };
 }
 
 function getEsmAst(root: ReturnType<typeof mdxProcessor.parse>) {

docs/migration-guide.md

@@ -145,6 +145,10 @@ Storybook now requires that MDX pages reference stories written in CSF, rather t
 
 You’ll also need to update your stories glob in `.storybook/main.js` to include the newly created .mdx and .stories.js files if it doesn’t already.
 
+#### Known limitations
+
+- The codemod does not remove the extracted stories from the `.stories.mdx` files. You will need to do this manually.
+
 **Note:** this migration supports the Storybook 6 ["CSF stories with MDX docs"](https://github.com/storybookjs/storybook/blob/6e19f0fe426d58f0f7981a42c3d0b0384fab49b1/code/addons/docs/docs/recipes.md#csf-stories-with-mdx-docs) recipe.
 
 ## Troubleshooting