From 12d3dc7fefad5a0411b94e30fa7fbd4c945df2cc Mon Sep 17 00:00:00 2001
From: Kasper Peulen <kasperpeulen@gmail.com>
Date: Fri, 13 Jan 2023 18:04:22 +0100
Subject: [PATCH] Upgrade deprecated type migration

---
 code/lib/codemod/package.json                 |   2 +
 .../transforms/__tests__/csf-2-to-3.test.ts   |  17 +-
 .../upgrade-deprecated-types.test.ts          | 170 ++++++++++++++++++
 code/lib/codemod/src/transforms/csf-2-to-3.ts |  24 +--
 .../transforms/upgrade-deprecated-types.ts    | 142 +++++++++++++++
 5 files changed, 342 insertions(+), 13 deletions(-)
 create mode 100644 code/lib/codemod/src/transforms/__tests__/upgrade-deprecated-types.test.ts
 create mode 100644 code/lib/codemod/src/transforms/upgrade-deprecated-types.ts

diff --git a/code/lib/codemod/package.json b/code/lib/codemod/package.json
index 7a480fc98a8c..ea3d17f6db33 100644
--- a/code/lib/codemod/package.json
+++ b/code/lib/codemod/package.json
@@ -34,6 +34,7 @@
     "./dist/transforms/storiesof-to-csf.js": "./dist/transforms/storiesof-to-csf.js",
     "./dist/transforms/update-addon-info.js": "./dist/transforms/update-addon-info.js",
     "./dist/transforms/update-organisation-name.js": "./dist/transforms/update-organisation-name.js",
+    "./dist/transforms/upgrade-deprecated-types.js": "./dist/transforms/upgrade-deprecated-types.js",
     "./dist/transforms/upgrade-hierarchy-separators.js": "./dist/transforms/upgrade-hierarchy-separators.js",
     "./package.json": "./package.json"
   },
@@ -81,6 +82,7 @@
       "./src/transforms/storiesof-to-csf.js",
       "./src/transforms/update-addon-info.js",
       "./src/transforms/update-organisation-name.js",
+      "./src/transforms/upgrade-deprecated-types.ts",
       "./src/transforms/upgrade-hierarchy-separators.js"
     ]
   },
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
index f7dcdf684a35..e20799ded383 100644
--- 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
@@ -234,7 +234,7 @@ describe('csf-2-to-3', () => {
         `)
       ).toThrowErrorMatchingInlineSnapshot(`
         This codemod does not support namespace imports for a @storybook/react package.
-                    Replace the namespace import with named imports and try again.
+        Replace the namespace import with named imports and try again.
         > 1 | import * as SB from '@storybook/react';
             | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
           2 | import { CatProps } from './Cat';
@@ -286,7 +286,7 @@ describe('csf-2-to-3', () => {
     it('should replace function exports with objects and update type', () => {
       expect(
         tsTransform(dedent`
-          import { Story, StoryFn, ComponentStory } from '@storybook/react';
+          import { Story, StoryFn, ComponentStory, ComponentStoryObj } from '@storybook/react';
 
           // some extra whitespace to test
 
@@ -315,9 +315,14 @@ describe('csf-2-to-3', () => {
             name: "Fluffy"
           };
           
+          export const G: ComponentStoryObj<typeof Cat> = {
+            args: {
+              name: 'Fluffy',
+            },
+          };
         `)
       ).toMatchInlineSnapshot(`
-        import { StoryObj, StoryFn, ComponentStory } from '@storybook/react';
+        import { StoryObj, StoryFn } from '@storybook/react';
 
         // some extra whitespace to test
 
@@ -351,6 +356,12 @@ describe('csf-2-to-3', () => {
             name: 'Fluffy',
           },
         };
+
+        export const G: StoryObj<typeof Cat> = {
+          args: {
+            name: 'Fluffy',
+          },
+        };
       `);
     });
   });
diff --git a/code/lib/codemod/src/transforms/__tests__/upgrade-deprecated-types.test.ts b/code/lib/codemod/src/transforms/__tests__/upgrade-deprecated-types.test.ts
new file mode 100644
index 000000000000..51a7640072c5
--- /dev/null
+++ b/code/lib/codemod/src/transforms/__tests__/upgrade-deprecated-types.test.ts
@@ -0,0 +1,170 @@
+import { describe, expect, it } from '@jest/globals';
+import { dedent } from 'ts-dedent';
+import type { API } from 'jscodeshift';
+import ansiRegex from 'ansi-regex';
+import _transform from '../upgrade-deprecated-types';
+
+expect.addSnapshotSerializer({
+  print: (val: any) => val,
+  test: () => true,
+});
+
+const tsTransform = (source: string) =>
+  _transform({ source, path: 'Component.stories.ts' }, {} as API, { parser: 'tsx' }).trim();
+
+describe('upgrade-deprecated-types', () => {
+  describe('typescript', () => {
+    it('upgrade regular imports', () => {
+      expect(
+        tsTransform(dedent`
+          import { Story, ComponentMeta, Meta, ComponentStory, ComponentStoryObj, ComponentStoryFn } from '@storybook/react';
+          import { Cat, CatProps } from './Cat';
+            
+          const meta = { title: 'Cat', component: Cat } satisfies ComponentMeta<typeof Cat>
+          const meta2: Meta<CatProps> = { title: 'Cat', component: Cat };
+          export default meta;
+          
+          export const A: ComponentStory<typeof Cat> = (args) => <Cat {...args} />;
+          export const B: any = (args) => <Button {...args} />;
+          export const C: ComponentStoryFn<typeof Cat> = (args) => <Cat {...args} />;
+          export const D: ComponentStoryObj<typeof Cat> = {
+            args: {
+              name: 'Fluffy',
+            },
+          };
+          export const E: Story<CatProps> = (args) => <Cat {...args} />;
+        `)
+      ).toMatchInlineSnapshot(`
+        import { StoryFn, Meta, StoryObj } from '@storybook/react';
+        import { Cat, CatProps } from './Cat';
+
+        const meta = { title: 'Cat', component: Cat } satisfies Meta<typeof Cat>;
+        const meta2: Meta<CatProps> = { title: 'Cat', component: Cat };
+        export default meta;
+
+        export const A: StoryFn<typeof Cat> = (args) => <Cat {...args} />;
+        export const B: any = (args) => <Button {...args} />;
+        export const C: StoryFn<typeof Cat> = (args) => <Cat {...args} />;
+        export const D: StoryObj<typeof Cat> = {
+          args: {
+            name: 'Fluffy',
+          },
+        };
+        export const E: StoryFn<CatProps> = (args) => <Cat {...args} />;
+      `);
+    });
+
+    it('upgrade imports with local names', () => {
+      expect(
+        tsTransform(dedent`
+          import { Story as Story_, ComponentMeta as ComponentMeta_, ComponentStory as Story__, ComponentStoryObj as ComponentStoryObj_, ComponentStoryFn as StoryFn_ } from '@storybook/react';
+          import { Cat } from './Cat';
+            
+          const meta = { title: 'Cat', component: Cat } satisfies ComponentMeta_<typeof Cat>
+          const meta2: ComponentMeta_<typeof Cat> = { title: 'Cat', component: Cat };
+          export default meta;
+          
+          export const A: Story__<typeof Cat> = (args) => <Cat {...args} />;
+          export const B: any = (args) => <Button {...args} />;
+          export const C: StoryFn_<typeof Cat> = (args) => <Cat {...args} />;
+          export const D: ComponentStoryObj_<typeof Cat> = {
+            args: {
+              name: 'Fluffy',
+            },
+          };
+          export const E: Story_<CatProps> = (args) => <Cat {...args} />;
+        `)
+      ).toMatchInlineSnapshot(`
+        import {
+          StoryFn as Story_,
+          Meta as ComponentMeta_,
+          StoryObj as ComponentStoryObj_,
+        } from '@storybook/react';
+        import { Cat } from './Cat';
+
+        const meta = { title: 'Cat', component: Cat } satisfies ComponentMeta_<typeof Cat>;
+        const meta2: ComponentMeta_<typeof Cat> = { title: 'Cat', component: Cat };
+        export default meta;
+
+        export const A: Story__<typeof Cat> = (args) => <Cat {...args} />;
+        export const B: any = (args) => <Button {...args} />;
+        export const C: StoryFn_<typeof Cat> = (args) => <Cat {...args} />;
+        export const D: ComponentStoryObj_<typeof Cat> = {
+          args: {
+            name: 'Fluffy',
+          },
+        };
+        export const E: Story_<CatProps> = (args) => <Cat {...args} />;
+      `);
+    });
+
+    it('upgrade imports with conflicting local names', () => {
+      expect.addSnapshotSerializer({
+        serialize: (value) => value.replace(ansiRegex(), ''),
+        test: () => true,
+      });
+
+      expect(() =>
+        tsTransform(dedent`
+          import { ComponentMeta as Meta, ComponentStory as StoryFn } from '@storybook/react';
+          import { Cat } from './Cat';
+            
+          const meta = { title: 'Cat', component: Cat } satisfies Meta<typeof Cat>
+          export default meta;
+          
+          export const A: StoryFn<typeof Cat> = (args) => <Cat {...args} />;
+         
+        `)
+      ).toThrowErrorMatchingInlineSnapshot(`
+        This codemod does not support local imports that are called the same as a storybook import.
+        Rename this local import and try again.
+        > 1 |  import { ComponentMeta as Meta, ComponentStory as StoryFn } from '@storybook/react';
+            |           ^^^^^^^^^^^^^^^^^^^^^
+          2 |  import { Cat } from './Cat';
+          3 |    
+          4 |  const meta = { title: 'Cat', component: Cat } satisfies Meta<typeof Cat>
+      `);
+    });
+
+    it('upgrade namespaces', () => {
+      expect(
+        tsTransform(dedent`
+          import * as SB from '@storybook/react';
+          import { Cat, CatProps } from './Cat';
+  
+          const meta = { title: 'Cat', component: Cat } satisfies SB.ComponentMeta<typeof Cat>;
+          const meta2: SB.ComponentMeta<typeof Cat> = { title: 'Cat', component: Cat };
+          export default meta;
+  
+          export const A: SB.ComponentStory<typeof Cat> = (args) => <Cat {...args} />;
+          export const B: any = (args) => <Button {...args} />;
+          export const C: SB.ComponentStoryFn<typeof Cat> = (args) => <Cat {...args} />;
+          export const D: SB.ComponentStoryObj<typeof Cat> = {
+            args: {
+              name: 'Fluffy',
+            },
+          };
+          export const E: SB.Story<CatProps> = (args) => <Cat {...args} />;
+          
+        `)
+      ).toMatchInlineSnapshot(`
+        import * as SB from '@storybook/react';
+        import { Cat, CatProps } from './Cat';
+
+        const meta = { title: 'Cat', component: Cat } satisfies SB.Meta<typeof Cat>;
+        const meta2: SB.Meta<typeof Cat> = { title: 'Cat', component: Cat };
+        export default meta;
+
+        export const A: SB.StoryFn<typeof Cat> = (args) => <Cat {...args} />;
+        export const B: any = (args) => <Button {...args} />;
+        export const C: SB.StoryFn<typeof Cat> = (args) => <Cat {...args} />;
+        export const D: SB.StoryObj<typeof Cat> = {
+          args: {
+            name: 'Fluffy',
+          },
+        };
+        export const E: SB.StoryFn<CatProps> = (args) => <Cat {...args} />;
+      `);
+    });
+  });
+});
diff --git a/code/lib/codemod/src/transforms/csf-2-to-3.ts b/code/lib/codemod/src/transforms/csf-2-to-3.ts
index c9de746299a5..59b4b16b8399 100644
--- a/code/lib/codemod/src/transforms/csf-2-to-3.ts
+++ b/code/lib/codemod/src/transforms/csf-2-to-3.ts
@@ -8,6 +8,7 @@ import type { API, FileInfo } from 'jscodeshift';
 import type { BabelFile, NodePath } from '@babel/core';
 import * as babel from '@babel/core';
 import * as recast from 'recast';
+import { upgradeDeprecatedTypes } from './upgrade-deprecated-types';
 
 const logger = console;
 
@@ -106,7 +107,14 @@ export default function transform(info: FileInfo, api: API, options: { parser?:
     return info.source;
   }
 
-  const importHelper = new StorybookImportHelper(csf, info);
+  // This allows for showing buildCodeFrameError messages
+  // @ts-expect-error File is not yet exposed, see https://github.com/babel/babel/issues/11350#issuecomment-644118606
+  const file: BabelFile = new babel.File(
+    { filename: info.path },
+    { code: info.source, ast: csf._ast }
+  );
+
+  const importHelper = new StorybookImportHelper(file, info);
 
   const objectExports: Record<string, t.Statement> = {};
   Object.entries(csf._storyExports).forEach(([key, decl]) => {
@@ -172,6 +180,8 @@ export default function transform(info: FileInfo, api: API, options: { parser?:
     return acc;
   }, []);
 
+  upgradeDeprecatedTypes(file);
+
   let output = recast.print(csf._ast, {}).code;
 
   try {
@@ -196,13 +206,7 @@ export default function transform(info: FileInfo, api: API, options: { parser?:
 }
 
 class StorybookImportHelper {
-  constructor(csf: CsfFile, info: FileInfo) {
-    // This allows for showing buildCodeFrameError messages
-    // @ts-expect-error File is not yet exposed, see https://github.com/babel/babel/issues/11350#issuecomment-644118606
-    const file: BabelFile = new babel.File(
-      { filename: info.path },
-      { code: info.source, ast: csf._ast }
-    );
+  constructor(file: BabelFile, info: FileInfo) {
     this.sbImportDeclarations = this.getAllSbImportDeclarations(file);
   }
 
@@ -218,8 +222,8 @@ class StorybookImportHelper {
         const isRendererImport = path.get('specifiers').some((specifier) => {
           if (specifier.isImportNamespaceSpecifier()) {
             throw path.buildCodeFrameError(
-              `This codemod does not support namespace imports for a ${path.node.source.value} package.
-            Replace the namespace import with named imports and try again.`
+              `This codemod does not support namespace imports for a ${path.node.source.value} package.\n` +
+                'Replace the namespace import with named imports and try again.'
             );
           }
           if (!specifier.isImportSpecifier()) return false;
diff --git a/code/lib/codemod/src/transforms/upgrade-deprecated-types.ts b/code/lib/codemod/src/transforms/upgrade-deprecated-types.ts
new file mode 100644
index 000000000000..6b76e75b5aa5
--- /dev/null
+++ b/code/lib/codemod/src/transforms/upgrade-deprecated-types.ts
@@ -0,0 +1,142 @@
+/* eslint-disable no-underscore-dangle */
+import prettier from 'prettier';
+import type { API, FileInfo } from 'jscodeshift';
+import type { BabelFile, NodePath } from '@babel/core';
+import * as babel from '@babel/core';
+import { loadCsf } from '@storybook/csf-tools';
+import * as recast from 'recast';
+import * as t from '@babel/types';
+
+const logger = console;
+
+const deprecatedTypes = [
+  'ComponentStory',
+  'ComponentStoryFn',
+  'ComponentStoryObj',
+  'ComponentMeta',
+  'Story',
+];
+
+function migrateType(oldType: string) {
+  if (oldType === 'Story' || oldType === 'ComponentStory') return 'StoryFn';
+  return oldType.replace('Component', '');
+}
+
+export default function transform(info: FileInfo, api: API, options: { parser?: string }) {
+  // TODO what do I need to with the title?
+  const fileNode = loadCsf(info.source, { makeTitle: (title) => title })._ast;
+  // @ts-expect-error File is not yet exposed, see https://github.com/babel/babel/issues/11350#issuecomment-644118606
+  const file: BabelFile = new babel.File(
+    { filename: info.path },
+    { code: info.source, ast: fileNode }
+  );
+
+  upgradeDeprecatedTypes(file);
+
+  let output = recast.print(file.path.node).code;
+
+  try {
+    const prettierConfig = prettier.resolveConfig.sync('.', { editorconfig: true }) || {
+      printWidth: 100,
+      tabWidth: 2,
+      bracketSpacing: true,
+      trailingComma: 'es5',
+      singleQuote: true,
+    };
+
+    output = prettier.format(output, { ...prettierConfig, filepath: info.path });
+  } catch (e) {
+    logger.log(`Failed applying prettier to ${info.path}.`);
+  }
+
+  return output;
+}
+
+export const parser = 'tsx';
+
+export function upgradeDeprecatedTypes(file: BabelFile) {
+  const importedNamespaces: Set<string> = new Set();
+  const typeReferencesToUpdate: Set<string> = new Set();
+  const existingImports: { name: string; isAlias: boolean; path: NodePath }[] = [];
+
+  file.path.traverse({
+    ImportDeclaration: (path) => {
+      existingImports.push(
+        ...path.get('specifiers').map((specifier) => ({
+          name: specifier.node.local.name,
+          isAlias: !(
+            specifier.isImportSpecifier() &&
+            t.isIdentifier(specifier.node.imported) &&
+            specifier.node.local.name === specifier.node.imported.name
+          ),
+          path: specifier,
+        }))
+      );
+
+      const source = path.node.source.value;
+      if (!source.startsWith('@storybook')) return;
+
+      path.get('specifiers').forEach((specifier) => {
+        if (specifier.isImportNamespaceSpecifier()) {
+          importedNamespaces.add(specifier.node.local.name);
+        }
+        if (!specifier.isImportSpecifier()) return;
+        const imported = specifier.get('imported');
+        if (!imported.isIdentifier()) return;
+
+        // if we find a deprecated import
+        if (deprecatedTypes.includes(imported.node.name)) {
+          // we don't have to rewrite type references for aliased imports
+          if (imported.node.name === specifier.node.local.name) {
+            typeReferencesToUpdate.add(specifier.node.local.name);
+          }
+
+          const newType = migrateType(imported.node.name);
+
+          // replace the deprecated import type when the new type isn't yet imported
+          // note that we don't replace the local name of the specifier
+          if (!existingImports.some((it) => it.name === newType)) {
+            imported.replaceWith(t.identifier(newType));
+            existingImports.push({ name: newType, isAlias: false, path: specifier });
+          } else {
+            // if the existing import has the same local name but is an alias we throw
+            // we could have imported the type with an alias, but seems to much effort
+            const existingImport = existingImports.find((it) => it.name === newType && it.isAlias);
+            if (existingImport) {
+              throw existingImport.path.buildCodeFrameError(
+                'This codemod does not support local imports that are called the same as a storybook import.\n' +
+                  'Rename this local import and try again.'
+              );
+            } else {
+              // if the type already exists, without being aliased
+              // we can safely remove the deprecated import now
+              specifier.remove();
+            }
+          }
+        }
+      });
+    },
+  });
+
+  file.path.traverse({
+    TSTypeReference: (path) => {
+      const typeName = path.get('typeName');
+      if (typeName.isIdentifier()) {
+        if (typeReferencesToUpdate.has(typeName.node.name)) {
+          typeName.replaceWith(t.identifier(migrateType(typeName.node.name)));
+        }
+      } else if (typeName.isTSQualifiedName()) {
+        // For example SB.StoryObj
+        const namespace = typeName.get('left');
+        if (namespace.isIdentifier()) {
+          if (importedNamespaces.has(namespace.node.name)) {
+            const right = typeName.get('right');
+            if (deprecatedTypes.includes(right.node.name)) {
+              right.replaceWith(t.identifier(migrateType(right.node.name)));
+            }
+          }
+        }
+      }
+    },
+  });
+}