CLI: In csf-factories codemod only remove types which are unused

code/lib/cli-storybook/src/codemod/helpers/remove-unused-types.test.ts

@@ -0,0 +1,126 @@
+import { describe, expect, it } from 'vitest';
+
+import { loadCsf, printCsf } from 'storybook/internal/csf-tools';
+
+import { dedent } from 'ts-dedent';
+
+import { getDiff } from '../../../../../core/src/core-server/utils/save-story/getDiff';
+import { removeUnusedTypes } from './remove-unused-types';
+
+expect.addSnapshotSerializer({
+  serialize: (val: any) => (typeof val === 'string' ? val : val.toString()),
+  test: () => true,
+});
+
+const unescape = (str: string) => str.replace(/\r\n/g, '\n');
+
+describe('removeUnusedTypes', () => {
+  const getTransformed = (source: string) => {
+    const csf = loadCsf(source, { makeTitle: () => 'FIXME' }).parse();
+    removeUnusedTypes(csf._ast.program, csf._ast);
+    return printCsf(csf).code;
+  };
+
+  it('should remove unused Storybook types', async () => {
+    const source = dedent`
+      import { Button } from './Button';
+      import { StoryFn, StoryObj } from '@storybook/react';
+      import type { Meta, MetaObj } from '@storybook/react';
+      import { type ComponentStory, type ComponentMeta } from '@storybook/react';
+
+      // unused types that should be removed
+      type UnusedAlias = Meta<typeof Button>;
+      type UnusedAlias2 = StoryObj<typeof Button>;
+      type UnusedAlias3 = ComponentStory<typeof Button>;
+      type UnusedAlias4 = ComponentMeta<typeof Button>;
+      type UnusedDeepType = {
+        foo: {
+          bar: {
+            story: StoryObj<typeof Button>;
+          }
+        }
+      };
+      interface UnusedInterface extends Meta {}
+      interface UnusedDeepInterface {
+        baz: {
+          qux: {
+            meta: Meta<typeof Button>;
+          }
+        }
+      };
+
+      export default { component: Button };
+    `;
+
+    const transformed = getTransformed(source);
+    expect(getDiff(unescape(source), unescape(transformed))).toMatchInlineSnapshot(`
+      import { Button } from './Button';
+
+      - import { StoryFn, StoryObj } from '@storybook/react';
+      - import type { Meta, MetaObj } from '@storybook/react';
+      - import { type ComponentStory, type ComponentMeta } from '@storybook/react';
+      - 
+
+
+      - // unused types that should be removed
+      - type UnusedAlias = Meta<typeof Button>;
+      - type UnusedAlias2 = StoryObj<typeof Button>;
+      - type UnusedAlias3 = ComponentStory<typeof Button>;
+      - type UnusedAlias4 = ComponentMeta<typeof Button>;
+      - type UnusedDeepType = {
+      -   foo: {
+      -     bar: {
+      -       story: StoryObj<typeof Button>;
+      -     }
+      -   }
+      - };
+      - interface UnusedInterface extends Meta {}
+      - interface UnusedDeepInterface {
+      -   baz: {
+      -     qux: {
+      -       meta: Meta<typeof Button>;
+      -     }
+      -   }
+      - };
+      - 
+      - 
+        export default { component: Button };
+    `);
+  });
+
+  it('should not remove used Storybook types', async () => {
+    const source = dedent`
+      // Nothing in this file should be removed or modified
+      import { StoryFn, StoryObj, ComponentStory, Meta, MetaObj, ComponentMeta } from '@storybook/react';
+      import { Button } from './Button';
+
+      type Alias = StoryFn<typeof Button>;
+      type Alias2 = Alias & { b: string };
+      type Story = StoryObj & { a: string };
+      type DeepType = {
+        foo: {
+          bar: {
+            story: ComponentStory<typeof Button>;
+          }
+        }
+      };
+      interface Interface extends Meta {}
+      interface DeepInterface {
+        baz: {
+          qux: {
+            meta: MetaObj<typeof Button>;
+          }
+        }
+      };
+      const X: ComponentMeta = {}
+
+      function foo(a: Story, c: DeepType, d: Interface, e: DeepInterface){}
+
+      export default {};
+    `;
+
+    const transformed = getTransformed(source);
+
+    expect(unescape(transformed)).toEqual(unescape(source));
+  });
+});

code/lib/cli-storybook/src/codemod/helpers/remove-unused-types.ts

@@ -0,0 +1,168 @@
+import { types as t, traverse } from 'storybook/internal/babel';
+
+import { cleanupTypeImports } from './csf-factories-utils';
+
+// Name of types that should be removed from the import list
+const typesDisallowList = [
+  'Story',
+  'StoryFn',
+  'StoryObj',
+  'Meta',
+  'MetaObj',
+  'ComponentStory',
+  'ComponentMeta',
+];
+
+const disallowedTypesSet = new Set(typesDisallowList);
+
+/**
+ * Remove unused Storybook-specific type aliases from the program.
+ *
+ * Conditions to remove a declared type/interface:
+ *
+ * - It is declared in the file,
+ * - It is not referenced anywhere in the file,
+ * - AND it (the declaration) references at least one Storybook type from typesDisallowList.
+ *
+ * This implementation performs a single traversal of `ast`. During traversal we:
+ *
+ * - Collect declared type names,
+ * - Record references to declared types (including handling references that appear before
+ *   declarations),
+ * - Detect per-declaration whether it references any disallowed Storybook type, and then perform a
+ *   single filter pass on program.body.
+ */
+export function removeUnusedTypes(programNode: t.Program, ast: t.File): void {
+  // Declared type/interface names seen in this file
+  const declaredTypes = new Set<string>();
+
+  // Names of declared types that are referenced somewhere in the file
+  const referencedTypes = new Set<string>();
+
+  // Temporary: identifier names seen before we encountered their declaration
+  // This lets us count forward references (identifier appears before type is declared).
+  const pendingIdentifierNames = new Set<string>();
+
+  // Names of type declarations that (somewhere in their AST) reference a disallowed Storybook type
+  const typeDeclReferencesDisallowed = new Set<string>();
+
+  traverse(ast, {
+    enter(path) {
+      const node = path.node;
+
+      // 1) When we encounter a type/interface declaration, register it.
+      if (path.isTSTypeAliasDeclaration() || path.isTSInterfaceDeclaration()) {
+        // These always have an `id` property that's an Identifier
+        const idNode = (node as t.TSTypeAliasDeclaration | t.TSInterfaceDeclaration).id;
+        const name = idNode && t.isIdentifier(idNode) ? idNode.name : undefined;
+        if (name) {
+          declaredTypes.add(name);
+
+          // If we previously saw identifiers with this name before the declaration,
+          // count them now as references (handles reference-before-declaration).
+          if (pendingIdentifierNames.has(name)) {
+            referencedTypes.add(name);
+          }
+        }
+
+        // No need to traverse into the id itself here; we still want to traverse the
+        // declaration body so that disallowed-type references inside are detected
+        // by the TSTypeReference/TSExpressionWithTypeArguments handlers below.
+        return;
+      }
+
+      // 2) Track identifier references to declared types.
+      if (path.isIdentifier()) {
+        const identifierNode = node as t.Identifier;
+        const name = identifierNode.name;
+
+        // Skip the identifier that *is* the declaration id itself:
+        // parent is TSTypeAliasDeclaration or TSInterfaceDeclaration and its id is this node
+        const parentPath = path.parentPath;
+        if (
+          parentPath &&
+          (parentPath.isTSTypeAliasDeclaration() || parentPath.isTSInterfaceDeclaration())
+        ) {
+          const parentIdNode = (
+            parentPath.node as t.TSTypeAliasDeclaration | t.TSInterfaceDeclaration
+          ).id;
+          if (parentIdNode === node) {
+            return;
+          }
+        }
+
+        // If we've already seen the declaration, mark as referenced.
+        if (declaredTypes.has(name)) {
+          referencedTypes.add(name);
+        } else {
+          // Otherwise record as pending — if the declaration appears later, we'll promote it.
+          pendingIdentifierNames.add(name);
+        }
+
+        return;
+      }
+
+      // 3) Detect references to disallowed Storybook types inside type declarations.
+      // If we find one, record which type declaration (owner) contains it.
+      if (path.isTSTypeReference()) {
+        const typeRefNode = node as t.TSTypeReference;
+        const typeNameNode = typeRefNode.typeName;
+
+        if (t.isIdentifier(typeNameNode) && disallowedTypesSet.has(typeNameNode.name)) {
+          // Find the nearest enclosing type declaration (alias or interface)
+          const owner = path.findParent(
+            (p) => p.isTSTypeAliasDeclaration() || p.isTSInterfaceDeclaration()
+          );
+          if (owner && (owner.isTSTypeAliasDeclaration() || owner.isTSInterfaceDeclaration())) {
+            const ownerId = (owner.node as t.TSTypeAliasDeclaration | t.TSInterfaceDeclaration).id;
+            const ownerName = t.isIdentifier(ownerId) ? ownerId.name : undefined;
+            if (ownerName) {
+              typeDeclReferencesDisallowed.add(ownerName);
+            }
+          }
+        }
+
+        return;
+      }
+
+      if (path.isTSExpressionWithTypeArguments()) {
+        const tsExprNode = node as t.TSExpressionWithTypeArguments;
+        const expr = tsExprNode.expression;
+        if (t.isIdentifier(expr) && disallowedTypesSet.has(expr.name)) {
+          const owner = path.findParent(
+            (p) => p.isTSTypeAliasDeclaration() || p.isTSInterfaceDeclaration()
+          );
+          if (owner && (owner.isTSTypeAliasDeclaration() || owner.isTSInterfaceDeclaration())) {
+            const ownerId = (owner.node as t.TSTypeAliasDeclaration | t.TSInterfaceDeclaration).id;
+            const ownerName = t.isIdentifier(ownerId) ? ownerId.name : undefined;
+            if (ownerName) {
+              typeDeclReferencesDisallowed.add(ownerName);
+            }
+          }
+        }
+        return;
+      }
+    },
+  });
+
+  // Final pass: remove unused declared types that reference disallowed types
+  programNode.body = programNode.body.filter((node) => {
+    if (t.isTSTypeAliasDeclaration(node) || t.isTSInterfaceDeclaration(node)) {
+      const name = node.id.name;
+
+      // If it's a declared type, unused, and references a disallowed Storybook type — remove it.
+      if (
+        declaredTypes.has(name) &&
+        !referencedTypes.has(name) &&
+        typeDeclReferencesDisallowed.has(name)
+      ) {
+        return false; // filter out (remove)
+      }
+    }
+
+    return true; // keep everything else
+  });
+
+  // Cleanup any now-unused Storybook type imports (keeps original API: pass array)
+  programNode.body = cleanupTypeImports(programNode, typesDisallowList);
+}

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

@@ -623,6 +623,84 @@ describe('stories codemod', () => {
       `);
     });
 
+    it('should preserve user-defined generic types', async () => {
+      const result = await transform(dedent`
+        import { Meta, StoryObj } from '@storybook/react';
+
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        type Data = Record<string, any>;
+        interface UnusedButShouldNotBeRemoved { name: string };
+        type UnusedAndShouldBeRemoved = Meta;
+
+        export default { title: 'Table' };
+
+        export const A = {
+          render: () => {
+            const data: Data[] = [];
+            return <Table data={data} />;
+          }
+        };
+      `);
+
+      expect(result).toContain('UnusedButShouldNotBeRemoved');
+      expect(result).not.toContain('UnusedAndShouldBeRemoved');
+
+      expect(result).toMatchInlineSnapshot(`
+        import preview from '#.storybook/preview';
+
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        type Data = Record<string, any>;
+        interface UnusedButShouldNotBeRemoved {
+          name: string;
+        }
+
+        const meta = preview.meta({
+          title: 'Table',
+        });
+
+        export const A = meta.story({
+          render: () => {
+            const data: Data[] = [];
+            return <Table data={data} />;
+          },
+        });
+      `);
+    });
+
+    it('should remove Storybook-specific type aliases but leave the ones that are actually used', async () => {
+      await expect(
+        transform(dedent`
+          import { Meta, StoryObj, ComponentStory, ComponentMeta } from '@storybook/react';
+          import { Button } from './Button';
+
+          type CustomMeta = Meta<typeof Button>;
+          type CustomStory = StoryObj<typeof Button>;
+          type LegacyStory = ComponentStory<typeof Button>;
+          type LegacyMeta = ComponentMeta<typeof Button>;
+          type ThisShouldNotBeRemoved = Meta<typeof Button>;
+          const something: ThisShouldNotBeRemoved = {};
+
+          export default { title: 'Button' };
+          export const A = {};
+        `)
+      ).resolves.toMatchInlineSnapshot(`
+        import { Meta } from '@storybook/react';
+
+        import preview from '#.storybook/preview';
+
+        import { Button } from './Button';
+
+        type ThisShouldNotBeRemoved = Meta<typeof Button>;
+        const something: ThisShouldNotBeRemoved = {};
+
+        const meta = preview.meta({
+          title: 'Button',
+        });
+
+        export const A = meta.story();
+      `);
+    });
+
     it.todo('should support non-conventional formats', async () => {
       const transformed = await transform(dedent`
         import { Meta, StoryObj as CSF3 } from '@storybook/react';