Merge pull request #16829 from storybookjs/16828-csf-tools-function-e…

…xports CSFFile: Fix function exports
storybookjs · Nov 30, 2021 · ee22b39 · ee22b39
2 parents 72735c6 + 4a1c8cc
commit ee22b39
Show file tree

Hide file tree

Showing 3 changed files with 69 additions and 34 deletions.
diff --git a/lib/codemod/src/transforms/csf-2-to-3.ts b/lib/codemod/src/transforms/csf-2-to-3.ts
@@ -104,36 +104,41 @@ function transform({ source }: { source: string }, api: any, options: { parser?:
       return t.objectProperty(t.identifier(_rename(annotation)), val as t.Expression);
     });
 
-    const { init, id } = decl;
-    // only replace arrow function expressions && template
-    // ignore no-arg stories without annotations
-    const template = getTemplateBindVariable(init);
-    if ((!t.isArrowFunctionExpression(init) && !template) || isSimpleCSFStory(init, annotations)) {
-      return;
-    }
+    if (t.isVariableDeclarator(decl)) {
+      const { init, id } = decl;
+      // only replace arrow function expressions && template
+      // ignore no-arg stories without annotations
+      const template = getTemplateBindVariable(init);
+      if (
+        (!t.isArrowFunctionExpression(init) && !template) ||
+        isSimpleCSFStory(init, annotations)
+      ) {
+        return;
+      }
 
-    // Remove the render function when we can hoist the template
-    // const Template = (args) => <Cat {...args} />;
-    // export const A = Template.bind({});
-    let storyFn = template && csf._templates[template];
-    if (!storyFn) storyFn = init;
-
-    const keyId = t.identifier(key);
-    // @ts-ignore
-    const { typeAnnotation } = id;
-    if (typeAnnotation) {
-      keyId.typeAnnotation = typeAnnotation;
-    }
+      // Remove the render function when we can hoist the template
+      // const Template = (args) => <Cat {...args} />;
+      // export const A = Template.bind({});
+      let storyFn = template && csf._templates[template];
+      if (!storyFn) storyFn = init;
+
+      const keyId = t.identifier(key);
+      // @ts-ignore
+      const { typeAnnotation } = id;
+      if (typeAnnotation) {
+        keyId.typeAnnotation = typeAnnotation;
+      }
 
-    const renderAnnotation = isReactGlobalRenderFn(csf, storyFn)
-      ? []
-      : [t.objectProperty(t.identifier('render'), storyFn)];
+      const renderAnnotation = isReactGlobalRenderFn(csf, storyFn)
+        ? []
+        : [t.objectProperty(t.identifier('render'), storyFn)];
 
-    objectExports[key] = t.exportNamedDeclaration(
-      t.variableDeclaration('const', [
-        t.variableDeclarator(keyId, t.objectExpression([...renderAnnotation, ...annotations])),
-      ])
-    );
+      objectExports[key] = t.exportNamedDeclaration(
+        t.variableDeclaration('const', [
+          t.variableDeclarator(keyId, t.objectExpression([...renderAnnotation, ...annotations])),
+        ])
+      );
+    }
   });
 
   const updatedBody = csf._ast.program.body.reduce((acc, stmt) => {

diff --git a/lib/csf-tools/src/CsfFile.test.ts b/lib/csf-tools/src/CsfFile.test.ts
@@ -405,6 +405,26 @@ describe('CsfFile', () => {
         )
       ).toThrow('CSF: unexpected storiesOf call');
     });
+
+    it('function exports', () => {
+      expect(
+        parse(
+          dedent`
+          export default { title: 'foo/bar' };
+          export function A() {}
+          export function B() {}
+      `
+        )
+      ).toMatchInlineSnapshot(`
+        meta:
+          title: foo/bar
+        stories:
+          - id: foo-bar--a
+            name: A
+          - id: foo-bar--b
+            name: B
+      `);
+    });
   });
 
   // NOTE: this does not have a public API, but we can still test it

diff --git a/lib/csf-tools/src/CsfFile.ts b/lib/csf-tools/src/CsfFile.ts
@@ -70,7 +70,7 @@ const formatLocation = (node: t.Node, fileName?: string) => {
   return `${fileName || ''} (line ${line}, col ${column})`.trim();
 };
 
-const isArgsStory = (init: t.Expression, parent: t.Node, csf: CsfFile) => {
+const isArgsStory = (init: t.Node, parent: t.Node, csf: CsfFile) => {
   let storyFn: t.Node = init;
   // export const Foo = Bar.bind({})
   if (t.isCallExpression(init)) {
@@ -98,6 +98,9 @@ const isArgsStory = (init: t.Expression, parent: t.Node, csf: CsfFile) => {
   if (t.isArrowFunctionExpression(storyFn)) {
     return storyFn.params.length > 0;
   }
+  if (t.isFunctionDeclaration(storyFn)) {
+    return storyFn.params.length > 0;
+  }
   return false;
 };
 
@@ -149,7 +152,7 @@ export class CsfFile {
 
   _metaAnnotations: Record<string, t.Node> = {};
 
-  _storyExports: Record<string, t.VariableDeclarator> = {};
+  _storyExports: Record<string, t.VariableDeclarator | t.FunctionDeclaration> = {};
 
   _storyAnnotations: Record<string, Record<string, t.Node>> = {};
 
@@ -231,12 +234,18 @@ export class CsfFile {
       },
       ExportNamedDeclaration: {
         enter({ node, parent }) {
+          let declarations;
           if (t.isVariableDeclaration(node.declaration)) {
+            declarations = node.declaration.declarations.filter((d) => t.isVariableDeclarator(d));
+          } else if (t.isFunctionDeclaration(node.declaration)) {
+            declarations = [node.declaration];
+          }
+          if (declarations) {
             // export const X = ...;
-            node.declaration.declarations.forEach((decl) => {
-              if (t.isVariableDeclarator(decl) && t.isIdentifier(decl.id)) {
+            declarations.forEach((decl: t.VariableDeclarator | t.FunctionDeclaration) => {
+              if (t.isIdentifier(decl.id)) {
                 const { name: exportName } = decl.id;
-                if (exportName === '__namedExportsOrder') {
+                if (exportName === '__namedExportsOrder' && t.isVariableDeclarator(decl)) {
                   self._namedExportsOrder = parseExportsOrder(decl.init);
                   return;
                 }
@@ -250,7 +259,7 @@ export class CsfFile {
                   self._storyAnnotations[exportName] = {};
                 }
                 let parameters;
-                if (t.isObjectExpression(decl.init)) {
+                if (t.isVariableDeclarator(decl) && t.isObjectExpression(decl.init)) {
                   let __isArgsStory = true; // assume default render is an args story
                   // CSF3 object export
                   decl.init.properties.forEach((p: t.ObjectProperty) => {
@@ -265,10 +274,11 @@ export class CsfFile {
                   });
                   parameters = { __isArgsStory };
                 } else {
+                  const fn = t.isVariableDeclarator(decl) ? decl.init : decl;
                   parameters = {
                     // __id: toId(self._meta.title, name),
                     // FIXME: Template.bind({});
-                    __isArgsStory: isArgsStory(decl.init, parent, self),
+                    __isArgsStory: isArgsStory(fn, parent, self),
                   };
                 }
                 self._stories[exportName] = {