codemod-barrel-files.ts

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 = [
  "src/services/my-service",
  "src/design-system",
  // Add more barrel files here
];
 
// Optional - List of paths where we want to drop the last segment
const DROP_LAST_SEGMENT_PATHS = [
  "src/design-system/components",
  "src/design-system/atoms",
  "src/design-system/molecules",
  "src/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,
  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,
        { 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;