Codemod to kill barrel file references

transform-barrel-file-imports.ts

/**
 * This script/codemod transforms imports from barrel files into direct imports.
 * It's designed to work with jscodeshift to modify TypeScript/JavaScript files.
 *
 * Features:
 * - Handles multiple barrel files
 * - Transforms both value and type imports
 * - Maintains correct relative paths
 * - Handles re-exports within barrel files
 *
 * Usage:
 * 1. Update the BARREL_IMPORTS array below with the paths of your barrel files.
 * 2. Update the DROP_LAST_SEGMENT_PATHS array if you want to drop the last
 *    segment of the path (for example, when importing from a folder with an
 *    index file).
 * 3. Run the script:
 *    npx jscodeshift -t ./codemods/transform-barrel-file-imports.ts --parser=ts ./src
 *
 * Options:
 * --dry: Use this flag to see what changes would be made without actually changing files
 * --run-in-band: Runs the script sequentially instead of in parallel (recommended for debugging)
 *
 * Example:
 * npx jscodeshift -t ./codemods/transform-barrel-file-imports.ts --parser=ts ./src --run-in-band --dry ./src > transform-log.txt
 *
 * Note: Always run with --dry first and review the changes before applying them to your codebase.
 *
 * This codemod was built with blood, sweat, and several AI prompts. If you need
 * to modify it, better ask for help to an AI (unless you're familiar with ASTs).
 */
import * as fs from "fs";
import * as path from "path";

import * as ts from "typescript";

import type { API, FileInfo, Options, Transform } from "jscodeshift";

// List of barrel files to process; in an ideal world this shouldn't be
// hardcoded, but it's a good starting point.
const BARREL_IMPORTS = [
  "services/my-service",
  "components/design-system",
  // Add more barrel files here
];

// List of paths where we want to drop the last segment
const DROP_LAST_SEGMENT_PATHS = [
  "components/design-system/components",
  "components/design-system/atoms",
  "components/design-system/molecules",
  "components/design-system/organisms"
];

// This map will store the real paths of all exported components, types, and enums
const exportedItemsMap = new Map<
  string,
  { path: string; kind: "value" | "type" }
>();

function getCompilerOptions(filePath: string): ts.CompilerOptions {
  const configPath = ts.findConfigFile(
    path.dirname(filePath),
    ts.sys.fileExists,
    "tsconfig.json"
  );
  if (!configPath) {
    throw new Error("Could not find a valid 'tsconfig.json'.");
  }

  const { config } = ts.readConfigFile(configPath, ts.sys.readFile);
  const { options } = ts.parseJsonConfigFileContent(
    config,
    ts.sys,
    path.dirname(configPath)
  );

  return options;
}

function resolveModule(
  importPath: string,
  containingFile: string
): string | null {
  const options = getCompilerOptions(containingFile);
  const moduleResolutionHost: ts.ModuleResolutionHost = {
    fileExists: ts.sys.fileExists,
    readFile: ts.sys.readFile,
    realpath: ts.sys.realpath,
    directoryExists: ts.sys.directoryExists,
    getCurrentDirectory: () => process.cwd(),
    getDirectories: ts.sys.getDirectories,
  };

  const resolved = ts.resolveModuleName(
    importPath,
    containingFile,
    options,
    moduleResolutionHost
  );

  return resolved.resolvedModule?.resolvedFileName || null;
}

function buildExportMap(filePath: string, visited = new Set<string>()) {
  if (visited.has(filePath)) return;
  visited.add(filePath);

  const fileContent = fs.readFileSync(filePath, "utf-8");
  const sourceFile = ts.createSourceFile(
    filePath,
    fileContent,
    ts.ScriptTarget.Latest,
    true
  );

  function visit(node: ts.Node) {
    if (ts.isExportDeclaration(node)) {
      if (node.exportClause && ts.isNamedExports(node.exportClause)) {
        node.exportClause.elements.forEach((element) => {
          const kind = element.isTypeOnly ? "type" : "value";
          if (node.moduleSpecifier) {
            const modulePath = (node.moduleSpecifier as ts.StringLiteral).text;
            const resolvedPath = resolveModule(modulePath, filePath);
            if (resolvedPath) {
              exportedItemsMap.set(element.name.text, {
                path: resolvedPath,
                kind,
              });
            }
          } else {
            exportedItemsMap.set(element.name.text, { path: filePath, kind });
          }
        });
      } else if (node.moduleSpecifier) {
        const modulePath = (node.moduleSpecifier as ts.StringLiteral).text;
        const resolvedPath = resolveModule(modulePath, filePath);
        if (resolvedPath) {
          buildExportMap(resolvedPath, visited);
        }
      }
    } else if (ts.isExportAssignment(node)) {
      exportedItemsMap.set("default", { path: filePath, kind: "value" });
    } else if (
      (ts.isFunctionDeclaration(node) ||
        ts.isClassDeclaration(node) ||
        ts.isVariableStatement(node) ||
        ts.isInterfaceDeclaration(node) ||
        ts.isTypeAliasDeclaration(node) ||
        ts.isEnumDeclaration(node)) &&
      node.modifiers?.some((m) => m.kind === ts.SyntaxKind.ExportKeyword)
    ) {
      if (
        ts.isFunctionDeclaration(node) ||
        ts.isClassDeclaration(node) ||
        ts.isInterfaceDeclaration(node) ||
        ts.isTypeAliasDeclaration(node) ||
        ts.isEnumDeclaration(node)
      ) {
        if (node.name) {
          const kind =
            ts.isInterfaceDeclaration(node) || ts.isTypeAliasDeclaration(node)
              ? "type"
              : "value";
          exportedItemsMap.set(node.name.text, { path: filePath, kind });
        }
      } else if (ts.isVariableStatement(node)) {
        node.declarationList.declarations.forEach((decl) => {
          if (ts.isIdentifier(decl.name)) {
            exportedItemsMap.set(decl.name.text, {
              path: filePath,
              kind: "value",
            });
          }
        });
      }
    }

    ts.forEachChild(node, visit);
  }

  visit(sourceFile);
}

const transform: Transform = (
  fileInfo: FileInfo,
  api: API,
  // eslint-disable-next-line @typescript-eslint/no-unused-vars
  options: Options
) => {
  const j = api.jscodeshift;
  const root = j(fileInfo.source);

  // Build the export map if it hasn't been built yet
  if (exportedItemsMap.size === 0) {
    BARREL_IMPORTS.forEach((barrelImport) => {
      const barrelPath = resolveModule(barrelImport, fileInfo.path);
      if (barrelPath) {
        buildExportMap(barrelPath);
      } else {
        console.warn(`Could not resolve barrel file: ${barrelImport}`);
      }
    });
  }

  let modified = false;

  root.find(j.ImportDeclaration).forEach((nodePath) => {
    const importPath = nodePath.node.source.value;

    const matchingBarrel = BARREL_IMPORTS.find(
      (barrel) => importPath === barrel || importPath.endsWith(`/${barrel}`)
    );

    if (matchingBarrel) {
      const newImports = new Map<
        string,
        // eslint-disable-next-line @typescript-eslint/no-explicit-any
        { valueSpecifiers: any[]; typeSpecifiers: any[] }
      >();

      nodePath.node.specifiers.forEach((specifier) => {
        if (specifier.type === "ImportSpecifier") {
          const itemName = specifier.imported.name;
          const localName = specifier.local.name;
          const exportedItem = exportedItemsMap.get(itemName);

          if (exportedItem) {
            // Get the path relative to the barrel file
            const barrelDir = path.dirname(
              resolveModule(matchingBarrel, fileInfo.path) || ""
            );
            let relativePath = path.relative(barrelDir, exportedItem.path);

            // If the relative path is empty, it means the export is from the barrel file itself
            if (relativePath === "") {
              relativePath = ".";
            }

            // Remove the file extension
            relativePath = relativePath.replace(/\.(ts|tsx|js|jsx)$/, "");

            // Ensure the path starts with the correct barrel import
            let newImportPath = path
              .join(matchingBarrel, relativePath)
              .replace(/\\/g, "/");

            // Check if we need to drop the last segment
            const shouldDropLastSegment = DROP_LAST_SEGMENT_PATHS.some(
              (dropPath) => newImportPath.startsWith(dropPath)
            );

            if (shouldDropLastSegment) {
              newImportPath = path.dirname(newImportPath);
            }

            if (!newImports.has(newImportPath)) {
              newImports.set(newImportPath, {
                valueSpecifiers: [],
                typeSpecifiers: [],
              });
            }

            const importGroup = newImports.get(newImportPath)!;
            const newSpecifier = j.importSpecifier(
              j.identifier(itemName),
              itemName !== localName ? j.identifier(localName) : null
            );

            if (
              exportedItem.kind === "type" ||
              specifier.importKind === "type"
            ) {
              importGroup.typeSpecifiers.push(newSpecifier);
            } else {
              importGroup.valueSpecifiers.push(newSpecifier);
            }
          } else {
            console.warn(`Could not find export information for ${itemName}`);
          }
        }
      });

      const newImportNodes = [...newImports.entries()].flatMap(
        ([importPath, { valueSpecifiers, typeSpecifiers }]) => {
          const imports = [];
          if (valueSpecifiers.length > 0) {
            imports.push(
              j.importDeclaration(valueSpecifiers, j.literal(importPath))
            );
          }
          if (typeSpecifiers.length > 0) {
            imports.push(
              j.importDeclaration(typeSpecifiers, j.literal(importPath), "type")
            );
          }
          return imports;
        }
      );

      if (newImportNodes.length > 0) {
        j(nodePath).replaceWith(newImportNodes);
        modified = true;
      }
    }
  });

  if (modified) {
    console.log(`Modified imports in ${fileInfo.path}`);
    return root.toSource();
  }

  return null;
};

export default transform;