Automigration: Fix wrap-require automigration for common js main.js files

code/lib/cli/src/automigrate/fixes/wrap-require.ts

@@ -61,7 +61,16 @@ export const wrapRequire: Fix<WrapRequireRunOptions> = {
           });
 
           if (getRequireWrapperName(mainConfig) === null) {
-            mainConfig.setImport(['dirname', 'join'], 'path');
+            if (
+              mainConfig.fileName.endsWith('.cjs') ||
+              mainConfig.fileName.endsWith('.cts') ||
+              mainConfig.fileName.endsWith('.cjsx') ||
+              mainConfig.fileName.endsWith('.ctsx')
+            ) {
+              mainConfig.setRequireImport(['dirname', 'join'], 'path');
+            } else {
+              mainConfig.setImport(['dirname', 'join'], 'path');
+            }
             mainConfig.setBodyDeclaration(
               getRequireWrapperAsCallExpression(result.isConfigTypescript)
             );

code/lib/csf-tools/src/ConfigFile.test.ts

@@ -1131,4 +1131,92 @@ describe('ConfigFile', () => {
       `);
     });
   });
+
+  describe('setRequireImport', () => {
+    it(`supports setting a default import for a field that does not exist`, () => {
+      const source = dedent`
+        const config: StorybookConfig = { };
+        export default config;
+      `;
+
+      const config = loadConfig(source).parse();
+      config.setRequireImport('path', 'path');
+
+      // eslint-disable-next-line no-underscore-dangle
+      const parsed = babelPrint(config._ast);
+
+      expect(parsed).toMatchInlineSnapshot(`
+        const path = require('path');
+        const config: StorybookConfig = { };
+        export default config;
+      `);
+    });
+
+    it(`supports setting a default import for a field that does exist`, () => {
+      const source = dedent`
+        const path = require('path');
+        const config: StorybookConfig = { };
+        export default config;
+      `;
+
+      const config = loadConfig(source).parse();
+      config.setRequireImport('path', 'path');
+
+      // eslint-disable-next-line no-underscore-dangle
+      const parsed = babelPrint(config._ast);
+
+      expect(parsed).toMatchInlineSnapshot(`
+        const path = require('path');
+        const config: StorybookConfig = { };
+        export default config;
+      `);
+    });
+
+    it(`supports setting a named import for a field that does not exist`, () => {
+      const source = dedent`
+        const config: StorybookConfig = { };
+        export default config;
+      `;
+
+      const config = loadConfig(source).parse();
+      config.setRequireImport(['dirname'], 'path');
+
+      // eslint-disable-next-line no-underscore-dangle
+      const parsed = babelPrint(config._ast);
+
+      expect(parsed).toMatchInlineSnapshot(`
+        const {
+          dirname,
+        } = require('path');
+
+        const config: StorybookConfig = { };
+        export default config;
+      `);
+    });
+
+    it(`supports setting a named import for a field where the source already exists`, () => {
+      const source = dedent`
+        const { dirname } = require('path');
+
+        const config: StorybookConfig = { };
+        export default config;
+      `;
+
+      const config = loadConfig(source).parse();
+      config.setRequireImport(['dirname', 'basename'], 'path');
+
+      // eslint-disable-next-line no-underscore-dangle
+      const parsed = babelPrint(config._ast);
+
+      expect(parsed).toMatchInlineSnapshot(`
+        const {
+          dirname,
+          basename,
+        } = require('path');
+
+        const config: StorybookConfig = { };
+        export default config;
+      `);
+    });
+  });
 });

code/lib/csf-tools/src/ConfigFile.ts

@@ -523,6 +523,107 @@ export class ConfigFile {
     this._ast.program.body.push(declaration);
   }
 
+  /**
+   * Import specifiers for a specific require import
+   * @param importSpecifiers - The import specifiers to set. If a string is passed in, a default import will be set. Otherwise, an array of named imports will be set
+   * @param fromImport - The module to import from
+   * @example
+   * // const { foo } = require('bar');
+   * setRequireImport(['foo'], 'bar');
+   *
+   * // const foo = require('bar');
+   * setRequireImport('foo', 'bar');
+   *
+   */
+  setRequireImport(importSpecifier: string[] | string, fromImport: string) {
+    const requireDeclaration = this._ast.program.body.find(
+      (node) =>
+        t.isVariableDeclaration(node) &&
+        node.declarations.length === 1 &&
+        t.isVariableDeclarator(node.declarations[0]) &&
+        t.isCallExpression(node.declarations[0].init) &&
+        t.isIdentifier(node.declarations[0].init.callee) &&
+        node.declarations[0].init.callee.name === 'require' &&
+        t.isStringLiteral(node.declarations[0].init.arguments[0]) &&
+        node.declarations[0].init.arguments[0].value === fromImport
+    ) as t.VariableDeclaration | undefined;
+
+    /**
+     * Returns true, when the given import declaration has the given import specifier
+     * @example
+     * // const { foo } = require('bar');
+     * hasImportSpecifier(declaration, 'foo');
+     */
+    const hasRequireSpecifier = (name: string) =>
+      t.isObjectPattern(requireDeclaration?.declarations[0].id) &&
+      requireDeclaration?.declarations[0].id.properties.find(
+        (specifier) =>
+          t.isObjectProperty(specifier) &&
+          t.isIdentifier(specifier.key) &&
+          specifier.key.name === name
+      );
+
+    /**
+     * Returns true, when the given import declaration has the given default import specifier
+     * @example
+     * // import foo from 'bar';
+     * hasImportSpecifier(declaration, 'foo');
+     */
+    const hasDefaultRequireSpecifier = (declaration: t.VariableDeclaration, name: string) =>
+      declaration.declarations.length === 1 &&
+      t.isVariableDeclarator(declaration.declarations[0]) &&
+      t.isIdentifier(declaration.declarations[0].id) &&
+      declaration.declarations[0].id.name === name;
+
+    // if the import specifier is a string, we're dealing with default imports
+    if (typeof importSpecifier === 'string') {
+      // If the import declaration with the given source exists
+      const addDefaultRequireSpecifier = () => {
+        this._ast.program.body.unshift(
+          t.variableDeclaration('const', [
+            t.variableDeclarator(
+              t.identifier(importSpecifier),
+              t.callExpression(t.identifier('require'), [t.stringLiteral(fromImport)])
+            ),
+          ])
+        );
+      };
+
+      if (requireDeclaration) {
+        if (!hasDefaultRequireSpecifier(requireDeclaration, importSpecifier)) {
+          // If the import declaration hasn't the specified default identifier, we add a new variable declaration
+          addDefaultRequireSpecifier();
+        }
+        // If the import declaration with the given source doesn't exist
+      } else {
+        // Add the import declaration to the top of the file
+        addDefaultRequireSpecifier();
+      }
+      // if the import specifier is an array, we're dealing with named imports
+    } else if (requireDeclaration) {
+      importSpecifier.forEach((specifier) => {
+        if (!hasRequireSpecifier(specifier)) {
+          (requireDeclaration.declarations[0].id as t.ObjectPattern).properties.push(
+            t.objectProperty(t.identifier(specifier), t.identifier(specifier), undefined, true)
+          );
+        }
+      });
+    } else {
+      this._ast.program.body.unshift(
+        t.variableDeclaration('const', [
+          t.variableDeclarator(
+            t.objectPattern(
+              importSpecifier.map((specifier) =>
+                t.objectProperty(t.identifier(specifier), t.identifier(specifier), undefined, true)
+              )
+            ),
+            t.callExpression(t.identifier('require'), [t.stringLiteral(fromImport)])
+          ),
+        ])
+      );
+    }
+  }
+
   /**
    * Set import specifiers for a given import statement.
    * @description Does not support setting type imports (yet)