Merge pull request #30352 from storybookjs/yann/fix-config-file-funct…

…ion-exports

Csf: Support named exports as functions

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

@@ -1466,5 +1466,16 @@ describe('ConfigFile', () => {
       expect(config._exportDecls['path']).toBe(undefined);
       expect(config._exports['path']).toBe(undefined);
     });
+
+    it('detects const and function export declarations', () => {
+      const source = dedent`
+        export function normalFunction() { };
+        export const value = ['@storybook/addon-essentials'];
+        export async function asyncFunction() { };
+        `;
+      const config = loadConfig(source).parse();
+
+      expect(Object.keys(config._exportDecls)).toHaveLength(3);
+    });
   });
 });

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

@@ -161,7 +161,7 @@ export class ConfigFile {
   // FIXME: this is a hack. this is only used in the case where the user is
   // modifying a named export that's a scalar. The _exports map is not suitable
   // for that. But rather than refactor the whole thing, we just use this as a stopgap.
-  _exportDecls: Record<string, t.VariableDeclarator> = {};
+  _exportDecls: Record<string, t.VariableDeclarator | t.FunctionDeclaration> = {};
 
   _exportsObject: t.ObjectExpression | undefined;
 
@@ -238,6 +238,13 @@ export class ConfigFile {
                 self._exportDecls[exportName] = decl;
               }
             });
+          } else if (t.isFunctionDeclaration(node.declaration)) {
+            // export function X() {...};
+            const decl = node.declaration;
+            if (t.isIdentifier(decl.id)) {
+              const { name: exportName } = decl.id;
+              self._exportDecls[exportName] = decl;
+            }
           } else if (node.specifiers) {
             // export { X };
             node.specifiers.forEach((spec) => {
@@ -380,7 +387,9 @@ export class ConfigFile {
       _updateExportNode(rest, expr, exportNode);
     } else if (exportNode && rest.length === 0 && this._exportDecls[path[0]]) {
       const decl = this._exportDecls[path[0]];
-      decl.init = _makeObjectExpression([], expr);
+      if (t.isVariableDeclarator(decl)) {
+        decl.init = _makeObjectExpression([], expr);
+      }
     } else if (this.hasDefaultExport) {
       // This means the main.js of the user has a default export that is not an object expression, therefore we can'types change the AST.
       throw new Error(

code/lib/cli-storybook/src/codemod/helpers/config-to-csf-factory.test.ts

@@ -58,10 +58,12 @@ describe('main/preview codemod: general parsing functionality', () => {
       });
     `);
   });
+
   it('should wrap defineMain call from const declared default export and default export mix', async () => {
     await expect(
       transform(dedent`
         export const tags = [];
+        export async function viteFinal(config) { return config };
         const config = {
           framework: '@storybook/react-vite',
         };
@@ -74,6 +76,9 @@ describe('main/preview codemod: general parsing functionality', () => {
       const config = {
         framework: '@storybook/react-vite',
         tags: [],
+        viteFinal: () => {
+          return config;
+        },
       };
 
       export default config;
@@ -82,16 +87,22 @@ describe('main/preview codemod: general parsing functionality', () => {
   it('should wrap defineMain call from named exports format', async () => {
     await expect(
       transform(dedent`
-        export const stories = ['../src/**/*.stories.@(js|jsx|ts|tsx)'];
+        export function stories() { return ['../src/**/*.stories.@(js|jsx|ts|tsx)'] };
         export const addons = ['@storybook/addon-essentials'];
+        export async function viteFinal(config) { return config };
         export const framework = '@storybook/react-vite';
       `)
     ).resolves.toMatchInlineSnapshot(`
       import { defineMain } from '@storybook/react-vite';
 
       export default defineMain({
-        stories: ['../src/**/*.stories.@(js|jsx|ts|tsx)'],
+        stories: () => {
+          return ['../src/**/*.stories.@(js|jsx|ts|tsx)'];
+        },
         addons: ['@storybook/addon-essentials'],
+        viteFinal: () => {
+          return config;
+        },
         framework: '@storybook/react-vite',
       });
     `);

code/lib/cli-storybook/src/codemod/helpers/config-to-csf-factory.ts

@@ -7,7 +7,11 @@ import picocolors from 'picocolors';
 
 import type { FileInfo } from '../../automigrate/codemod';
 import { logger } from '../csf-factories';
-import { cleanupTypeImports } from './csf-factories-utils';
+import {
+  cleanupTypeImports,
+  getConfigProperties,
+  removeExportDeclarations,
+} from './csf-factories-utils';
 
 export async function configToCsfFactory(
   info: FileInfo,
@@ -48,23 +52,10 @@ export async function configToCsfFactory(
   if (config._exportsObject && hasNamedExports) {
     const exportDecls = config._exportDecls;
 
-    for (const [name, decl] of Object.entries(exportDecls)) {
-      if (decl.init) {
-        config._exportsObject.properties.push(t.objectProperty(t.identifier(name), decl.init));
-      }
-    }
+    const defineConfigProps = getConfigProperties(exportDecls);
+    config._exportsObject.properties.push(...defineConfigProps);
 
-    programNode.body = programNode.body.filter((node) => {
-      if (t.isExportNamedDeclaration(node) && node.declaration) {
-        if (t.isVariableDeclaration(node.declaration)) {
-          node.declaration.declarations = node.declaration.declarations.filter(
-            (decl) => t.isIdentifier(decl.id) && !exportDecls[decl.id.name]
-          );
-          return node.declaration.declarations.length > 0;
-        }
-      }
-      return true;
-    });
+    programNode.body = removeExportDeclarations(programNode, exportDecls);
   } else if (config._exportsObject) {
     /**
      * Scenario 2: Default exports
@@ -124,32 +115,15 @@ export async function configToCsfFactory(
      * Transform into: export default defineMain({ foo: {}, bar: '' });
      */
     const exportDecls = config._exportDecls;
-    const defineConfigProps = [];
-
-    // Collect properties from named exports
-    for (const [name, decl] of Object.entries(exportDecls)) {
-      if (decl.init) {
-        defineConfigProps.push(t.objectProperty(t.identifier(name), decl.init));
-      }
-    }
+    const defineConfigProps = getConfigProperties(exportDecls);
 
     // Construct the `define` call
     const defineConfigCall = t.callExpression(t.identifier(methodName), [
       t.objectExpression(defineConfigProps),
     ]);
 
     // Remove all related named exports
-    programNode.body = programNode.body.filter((node) => {
-      if (t.isExportNamedDeclaration(node) && node.declaration) {
-        if (t.isVariableDeclaration(node.declaration)) {
-          node.declaration.declarations = node.declaration.declarations.filter(
-            (decl) => t.isIdentifier(decl.id) && !exportDecls[decl.id.name]
-          );
-          return node.declaration.declarations.length > 0;
-        }
-      }
-      return true;
-    });
+    programNode.body = removeExportDeclarations(programNode, exportDecls);
 
     // Add the new export default declaration
     programNode.body.push(t.exportDefaultDeclaration(defineConfigCall));

code/lib/cli-storybook/src/codemod/helpers/csf-factories-utils.ts

@@ -31,3 +31,48 @@ export function cleanupTypeImports(programNode: t.Program, disallowList: string[
     // error TS4058: Return type of exported function has or is using name 'BlockStatement' from external module "/code/core/dist/babel/index" but cannot be named
   }) as any;
 }
+
+export function removeExportDeclarations(
+  programNode: t.Program,
+  exportDecls: Record<string, t.VariableDeclarator | t.FunctionDeclaration>
+) {
+  return programNode.body.filter((node) => {
+    if (t.isExportNamedDeclaration(node) && node.declaration) {
+      if (t.isVariableDeclaration(node.declaration)) {
+        // Handle variable declarations
+        node.declaration.declarations = node.declaration.declarations.filter(
+          (decl) => t.isIdentifier(decl.id) && !exportDecls[decl.id.name]
+        );
+        return node.declaration.declarations.length > 0;
+      } else if (t.isFunctionDeclaration(node.declaration)) {
+        // Handle function declarations
+        const funcDecl = node.declaration;
+        return t.isIdentifier(funcDecl.id) && !exportDecls[funcDecl.id.name];
+      }
+    }
+    return true;
+    // @TODO adding any for now, unsure how to fix the following error:
+    // error TS4058: Return type of exported function has or is using name 'ObjectProperty' from external module "/tmp/storybook/code/core/dist/babel/index" but cannot be named.
+  }) as any;
+}
+
+export function getConfigProperties(
+  exportDecls: Record<string, t.VariableDeclarator | t.FunctionDeclaration>
+) {
+  const properties = [];
+
+  // Collect properties from named exports
+  for (const [name, decl] of Object.entries(exportDecls)) {
+    if (t.isVariableDeclarator(decl) && decl.init) {
+      properties.push(t.objectProperty(t.identifier(name), decl.init));
+    } else if (t.isFunctionDeclaration(decl)) {
+      properties.push(
+        t.objectProperty(t.identifier(name), t.arrowFunctionExpression([], decl.body))
+      );
+    }
+  }
+
+  // @TODO adding any for now, unsure how to fix the following error:
+  // error TS4058: Return type of exported function has or is using name 'ObjectProperty' from external module "/tmp/storybook/code/core/dist/babel/index" but cannot be named.
+  return properties as any;
+}