Repository: storybookjs/storybook
File: code/lib/cli-storybook/src/codemod/helpers/remove-unused-types.ts (TypeScript)

Code before:
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',
];

// Helper to detect if a type node references (directly) any type in typesDisallowList
function isReferencingDisallowedType(
  typeNode: t.TSTypeAliasDeclaration | t.TSInterfaceDeclaration
): boolean {
  let found = false;
  traverse(typeNode, {
    noScope: true,
    TSTypeReference(path) {
      if (
        t.isIdentifier(path.node.typeName) &&
        typesDisallowList.includes(path.node.typeName.name)
      ) {
        found = true;
        path.stop();
      }
    },
    TSExpressionWithTypeArguments(path) {
      if (
        t.isIdentifier(path.node.expression) &&
        typesDisallowList.includes(path.node.expression.name)
      ) {
        found = true;
        path.stop();
      }
    },
  });
  return found;
}

/**
 * Remove unused Storybook-specific type aliases from the program
 *
 * Only removes types that:
 *
 * 1. Reference Storybook types (Meta, Story, StoryObj, etc.)
 * 2. Are not actually used anywhere in the code
 */
export function removeUnusedTypes(programNode: t.Program, ast: t.File): void {
  // Find all type and interface declarations in the AST
  const declaredTypes = new Map<string, t.TSTypeAliasDeclaration | t.TSInterfaceDeclaration>();

  traverse(ast, {
    TSTypeAliasDeclaration(path) {
      const name = path.node.id.name;
      declaredTypes.set(name, path.node);
    },
    TSInterfaceDeclaration(path) {
      const name = path.node.id.name;
      declaredTypes.set(name, path.node);
    },
  });

  const declaredTypeNames = new Set(declaredTypes.keys());
  const referencedTypes = new Set<string>();

  traverse(ast, {
    Identifier(path) {
      const { node, parent } = path;

      // Only track as "used" if it's not the id of a declaration and refers to a declared type/interface
      if (
        declaredTypeNames.has(node.name) &&
        !(
          (t.isTSTypeAliasDeclaration(parent) && parent.id === node) ||
          (t.isTSInterfaceDeclaration(parent) && parent.id === node)
        )
      ) {
        referencedTypes.add(node.name);
      }
    },
  });

  // First: Removes (only unused) types that reference Storybook types
  programNode.body = programNode.body.filter((node) => {
    if (
      (t.isTSTypeAliasDeclaration(node) || t.isTSInterfaceDeclaration(node)) &&
      declaredTypes.has(node.id.name) &&
      !referencedTypes.has(node.id.name)
    ) {
      // Remove ONLY IF this type/interface refers to any from typesDisallowList
      if (isReferencingDisallowedType(node)) {
        return false;
      }
    }
    return true;
  });

  // Second: Remove (only unused) type imports – now inferred – from @storybook/* packages
  programNode.body = cleanupTypeImports(programNode, typesDisallowList);
}


Review comment (suggestion):
This can be done with a single iteration/traversal:

```suggestion
export function removeUnusedTypes(programNode: t.Program, ast: t.File): void {
  const declaredTypes = new Map<string, t.TSTypeAliasDeclaration | t.TSInterfaceDeclaration>();
  const declaredTypeNames = new Set<string>();
  const referencedTypes = new Set<string>();

  traverse(ast, {
    enter(path) {
      const node = path.node;

      // 1. Register declarations
      if (t.isTSTypeAliasDeclaration(node) || t.isTSInterfaceDeclaration(node)) {
        const name = node.id.name;
        declaredTypes.set(name, node);
        declaredTypeNames.add(name);
        return;
      }

      // 2. Track references to declared type/interface names
      if (t.isIdentifier(node)) {
        const parent = path.parent;

        // skip the identifier that *declares* the type/interface
        const isDeclId =
          (t.isTSTypeAliasDeclaration(parent) && parent.id === node) ||
          (t.isTSInterfaceDeclaration(parent) && parent.id === node);

        if (!isDeclId && declaredTypeNames.has(node.name)) {
          referencedTypes.add(node.name);
        }
      }
    },
  });

  // 3. Filter out unused declared types referencing disallowed types
  programNode.body = programNode.body.filter((node) => {
    if (
      (t.isTSTypeAliasDeclaration(node) || t.isTSInterfaceDeclaration(node)) &&
      declaredTypes.has(node.id.name) &&
      !referencedTypes.has(node.id.name)
    ) {
      // Remove ONLY IF this unused type references a disallowed type
      if (isReferencingDisallowedType(node)) {
        return false;
      }
    }
    return true;
  });

  // 4. Cleanup unused Storybook type imports
  programNode.body = cleanupTypeImports(programNode, typesDisallowList);
}
```

Code after:
        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);
}