CSF-2-3 automigration scripts ignores Typescript types

storybookjs · Dec 14, 2022 · 4b27a40 · 4b27a40
1 parent 3b6d967
commit 4b27a40
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 12 deletions.
diff --git a/code/lib/codemod/src/transforms/__tests__/csf-2-to-3.test.ts b/code/lib/codemod/src/transforms/__tests__/csf-2-to-3.test.ts
@@ -228,22 +228,40 @@ describe('csf-2-to-3', () => {
     it('should replace function exports with objects', () => {
       expect(
         tsTransform(dedent`
-          export default { title: 'Cat' } as Meta<CatProps>;
+          export default { title: 'Cat', component: Cat } as Meta<CatProps>;
           export const A: Story<CatProps> = (args) => <Cat {...args} />;
+          A.args = { name: "Fluffy" };
           export const B: any = (args) => <Button {...args} />;
           export const C: Story<CatProps> = () => <Cat />;
+          export const D: StoryFn<CatProps> = (args) => <Cat {...args} />;
+          D.args = { name: "Fluffy" };
+          export const E: ComponentStory<Cat> = (args) => <Cat {...args} />;
+          E.args = { name: "Fluffy" };
         `)
       ).toMatchInlineSnapshot(`
         export default {
           title: 'Cat',
+          component: Cat,
         } as Meta<CatProps>;
-        export const A: Story<CatProps> = {
-          render: (args) => <Cat {...args} />,
+        export const A: StoryObj<CatProps> = {
+          args: {
+            name: 'Fluffy',
+          },
         };
         export const B: any = {
           render: (args) => <Button {...args} />,
         };
-        export const C: Story<CatProps> = () => <Cat />;
+        export const C: StoryFn<CatProps> = () => <Cat />;
+        export const D: StoryObj<CatProps> = {
+          args: {
+            name: 'Fluffy',
+          },
+        };
+        export const E: StoryObj<Cat> = {
+          args: {
+            name: 'Fluffy',
+          },
+        };
       `);
     });
   });

diff --git a/code/lib/codemod/src/transforms/csf-2-to-3.ts b/code/lib/codemod/src/transforms/csf-2-to-3.ts
@@ -1,6 +1,7 @@
 /* eslint-disable no-underscore-dangle */
 import prettier from 'prettier';
 import * as t from '@babel/types';
+import { isIdentifier, isTSTypeAnnotation, isTSTypeReference } from '@babel/types';
 import type { CsfFile } from '@storybook/csf-tools';
 import { formatCsf, loadCsf } from '@storybook/csf-tools';
 import { jscodeshiftToPrettierParser } from '../lib/utils';
@@ -89,6 +90,26 @@ const isReactGlobalRenderFn = (csf: CsfFile, storyFn: t.Expression) => {
 const isSimpleCSFStory = (init: t.Expression, annotations: t.ObjectProperty[]) =>
   annotations.length === 0 && t.isArrowFunctionExpression(init) && init.params.length === 0;
 
+function updateTypeTo(id: t.LVal, type: string): t.LVal {
+  if (
+    isIdentifier(id) &&
+    isTSTypeAnnotation(id.typeAnnotation) &&
+    isTSTypeReference(id.typeAnnotation.typeAnnotation) &&
+    isIdentifier(id.typeAnnotation.typeAnnotation.typeName)
+  ) {
+    const { name } = id.typeAnnotation.typeAnnotation.typeName;
+    if (name === 'Story' || name === 'StoryFn' || name === 'ComponentStory') {
+      return {
+        ...id,
+        typeAnnotation: t.tsTypeAnnotation(
+          t.tsTypeReference(t.identifier(type), id.typeAnnotation.typeAnnotation.typeParameters)
+        ),
+      };
+    }
+  }
+  return { ...id };
+}
+
 function transform({ source }: { source: string }, api: any, options: { parser?: string }) {
   const makeTitle = (userTitle?: string) => {
     return userTitle || 'FIXME';
@@ -117,6 +138,9 @@ function transform({ source }: { source: string }, api: any, options: { parser?:
         (!t.isArrowFunctionExpression(init) && !template) ||
         isSimpleCSFStory(init, annotations)
       ) {
+        objectExports[key] = t.exportNamedDeclaration(
+          t.variableDeclaration('const', [t.variableDeclarator(updateTypeTo(id, 'StoryFn'), init)])
+        );
         return;
       }
 
@@ -128,20 +152,16 @@ function transform({ source }: { source: string }, api: any, options: { parser?:
         storyFn = init;
       }
 
-      const keyId = t.identifier(key);
-      // @ts-expect-error (Converted from ts-ignore)
-      const { typeAnnotation } = id;
-      if (typeAnnotation) {
-        keyId.typeAnnotation = typeAnnotation;
-      }
-
       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])),
+          t.variableDeclarator(
+            updateTypeTo(id, 'StoryObj'),
+            t.objectExpression([...renderAnnotation, ...annotations])
+          ),
         ])
       );
     }