feat: add replaceDefaultWithNamedImport util and refactor other utils

codemods/invariant/index.js

@@ -1,5 +1,5 @@
 import { ts } from '@ast-grep/napi';
-import { findNamedDefaultImport } from '../shared-ast-grep.js';
+import { replaceDefaultImport } from '../shared-ast-grep.js';
 
 /**
  * @typedef {import('../../types.js').Codemod} Codemod
@@ -17,32 +17,12 @@ export default function (options) {
 		transform: ({ file }) => {
 			const ast = ts.parse(file.source);
 			const root = ast.root();
-			const edits = [];
 
-			const imports = findNamedDefaultImport(root, 'invariant');
-
-			for (const imp of imports) {
-				const nameMatch = imp.getMatch('NAME');
-				if (nameMatch) {
-					const name = nameMatch.text();
-					const source = imp.text();
-					const quoteType = source.includes('"') ? '"' : "'";
-
-					if (imp.text().startsWith('import')) {
-						edits.push(
-							imp.replace(
-								`import ${name} from ${quoteType}tiny-invariant${quoteType};`,
-							),
-						);
-					} else {
-						edits.push(
-							imp.replace(
-								`const ${name} = require(${quoteType}tiny-invariant${quoteType});`,
-							),
-						);
-					}
-				}
-			}
+			const { edits } = replaceDefaultImport(
+				root,
+				'invariant',
+				'tiny-invariant',
+			);
 
 			return root.commitEdits(edits);
 		},

codemods/iterate-iterator/index.js

@@ -1,5 +1,5 @@
 import { ts } from '@ast-grep/napi';
-import { findNamedDefaultImport } from '../shared-ast-grep.js';
+import { removeImport } from '../shared-ast-grep.js';
 
 /**
  * @typedef {import('../../types.js').Codemod} Codemod
@@ -17,20 +17,10 @@ export default function (options) {
 		transform: ({ file }) => {
 			const ast = ts.parse(file.source);
 			const root = ast.root();
-			const edits = [];
-			const importNames = new Set();
 
-			const imports = findNamedDefaultImport(root, 'iterate-iterator');
+			const { edits, localNames } = removeImport(root, 'iterate-iterator');
 
-			for (const imp of imports) {
-				const nameMatch = imp.getMatch('NAME');
-				if (nameMatch) {
-					importNames.add(nameMatch.text());
-				}
-				edits.push(imp.replace(''));
-			}
-
-			for (const importName of importNames) {
+			for (const importName of localNames) {
 				const singleArgCalls = root.findAll({
 					rule: {
 						pattern: {

codemods/shared-ast-grep.js

@@ -91,13 +91,13 @@ export function removeImport(root, moduleName) {
 }
 
 /**
- * Find all named imports from a specific module in the AST.
+ * Find all default imports/requires for a package and extract common metadata.
  *
  * @param {SgNode} root - The root of the AST.
- * @param {string} moduleName - The name of the module to find imports from.
- * @returns {SgNode[]} - An array of matched import nodes.
+ * @param {string} moduleName - The package to find imports for.
+ * @returns {{ imports: SgNode[], localNames: string[], quoteType: string }}
  */
-export function findNamedDefaultImport(root, moduleName) {
+function findNamedDefaultImports(root, moduleName) {
 	const imports = root.findAll({
 		rule: {
 			any: [
@@ -135,7 +135,19 @@ export function findNamedDefaultImport(root, moduleName) {
 		},
 	});
 
-	return imports;
+	/** @type {string[]} */
+	const localNames = [];
+	let quoteType = "'";
+
+	for (const imp of imports) {
+		const nameMatch = imp.getMatch('NAME');
+		if (!nameMatch) continue;
+		localNames.push(nameMatch.text());
+		const impText = imp.text();
+		if (impText.includes('"')) quoteType = '"';
+	}
+
+	return { imports, localNames, quoteType };
 }
 
 /**
@@ -146,7 +158,7 @@ export function findNamedDefaultImport(root, moduleName) {
  * @returns {{ imports: SgNode[], identifierName: string | null }}
  */
 export function findDefaultImportIdentifier(root, moduleName) {
-	const imports = findNamedDefaultImport(root, moduleName);
+	const { imports } = findNamedDefaultImports(root, moduleName);
 	const identifierName = imports[0]?.getMatch('NAME')?.text() ?? null;
 	return { imports, identifierName };
 }
@@ -230,11 +242,12 @@ export function computePolyfillMethodCallReplacementEdits(
 
 /**
  * Replace a default import/require of one package with another.
+ * If toIdentifier is not provided, the original local name is preserved.
  *
  * @param {SgNode} root - The root of the AST.
  * @param {string} fromPackage - The package being replaced
  * @param {string} toPackage - The new package specifier
- * @param {string} toIdentifier - The local name to use in the replacement
+ * @param {string} [toIdentifier] - The local name to use in the replacement (defaults to original name)
  * @returns {{ edits: Edit[], localNames: string[], quoteType: string }}
  */
 export function replaceDefaultImport(
@@ -243,58 +256,78 @@ export function replaceDefaultImport(
 	toPackage,
 	toIdentifier,
 ) {
-	const imports = root.findAll({
-		rule: {
-			any: [
-				{
-					pattern: {
-						context: `import $NAME from '${fromPackage}'`,
-						strictness: 'relaxed',
-					},
-				},
-				{
-					pattern: {
-						context: `const $NAME = require('${fromPackage}')`,
-						strictness: 'relaxed',
-					},
-				},
-				{
-					pattern: {
-						context: `var $NAME = require('${fromPackage}')`,
-						strictness: 'relaxed',
-					},
-				},
-			],
-		},
-	});
+	const { imports, localNames, quoteType } = findNamedDefaultImports(
+		root,
+		fromPackage,
+	);
 
 	/** @type {Edit[]} */
 	const edits = [];
-	/** @type {string[]} */
-	const localNames = [];
-	let quoteType = "'";
 
 	for (const imp of imports) {
 		const nameMatch = imp.getMatch('NAME');
 		if (!nameMatch) continue;
 
-		localNames.push(nameMatch.text());
+		const identifier = toIdentifier || nameMatch.text();
+		const isCommonJS = imp.find('require($SOURCE)') !== null;
 
-		const impText = imp.text();
-		if (impText.includes('"')) quoteType = '"';
+		if (isCommonJS) {
+			edits.push(
+				imp.replace(
+					`const ${identifier} = require(${quoteType}${toPackage}${quoteType});`,
+				),
+			);
+		} else {
+			edits.push(
+				imp.replace(
+					`import ${identifier} from ${quoteType}${toPackage}${quoteType};`,
+				),
+			);
+		}
+	}
+
+	return { edits, localNames, quoteType };
+}
+
+/**
+ * Replace a default import/require of one package with a named import from another package.
+ *
+ * @param {SgNode} root - The root of the AST.
+ * @param {string} fromPackage - The package being replaced
+ * @param {string} toPackage - The new package specifier
+ * @param {string} namedImport - The named import to use (e.g., 'stripVTControlCharacters')
+ * @returns {{ edits: Edit[], localNames: string[], quoteType: string }}
+ */
+export function replaceDefaultWithNamedImport(
+	root,
+	fromPackage,
+	toPackage,
+	namedImport,
+) {
+	const { imports, localNames, quoteType } = findNamedDefaultImports(
+		root,
+		fromPackage,
+	);
+
+	/** @type {Edit[]} */
+	const edits = [];
+
+	for (const imp of imports) {
+		const nameMatch = imp.getMatch('NAME');
+		if (!nameMatch) continue;
 
 		const isCommonJS = imp.find('require($SOURCE)') !== null;
 
 		if (isCommonJS) {
 			edits.push(
 				imp.replace(
-					`const ${toIdentifier} = require(${quoteType}${toPackage}${quoteType});`,
+					`const { ${namedImport} } = require(${quoteType}${toPackage}${quoteType});`,
 				),
 			);
 		} else {
 			edits.push(
 				imp.replace(
-					`import ${toIdentifier} from ${quoteType}${toPackage}${quoteType};`,
+					`import { ${namedImport} } from ${quoteType}${toPackage}${quoteType};`,
 				),
 			);
 		}

codemods/strip-ansi/bun/index.js

@@ -1,5 +1,8 @@
 import { ts } from '@ast-grep/napi';
-import { findNamedDefaultImport } from '../../shared-ast-grep.js';
+import {
+	computeSimpleCallReplacementEdits,
+	replaceDefaultWithNamedImport,
+} from '../../shared-ast-grep.js';
 
 /**
  * @typedef {import('../../../types.js').Codemod} Codemod
@@ -17,52 +20,21 @@ export default function (options) {
 		transform: ({ file }) => {
 			const ast = ts.parse(file.source);
 			const root = ast.root();
-			const edits = [];
-			const importNames = new Set();
 
-			const imports = findNamedDefaultImport(root, 'strip-ansi');
-
-			for (const imp of imports) {
-				const nameMatch = imp.getMatch('NAME');
-				if (nameMatch) {
-					const name = nameMatch.text();
-					const source = imp.text();
-					const quoteType = source.includes('"') ? '"' : "'";
-
-					importNames.add(name);
-
-					if (imp.text().startsWith('import')) {
-						edits.push(
-							imp.replace(
-								`import { stripANSI } from ${quoteType}bun${quoteType};`,
-							),
-						);
-					} else {
-						edits.push(
-							imp.replace(
-								`const { stripANSI } = require(${quoteType}bun${quoteType});`,
-							),
-						);
-					}
-				}
-			}
-
-			for (const importName of importNames) {
-				const functionCalls = root.findAll({
-					rule: {
-						pattern: {
-							context: `${importName}($VALUE)`,
-							strictness: 'relaxed',
-						},
-					},
-				});
-
-				for (const call of functionCalls) {
-					const valueMatch = call.getMatch('VALUE');
-					if (valueMatch) {
-						edits.push(call.replace(`stripANSI(${valueMatch.text()})`));
-					}
-				}
+			const { edits, localNames } = replaceDefaultWithNamedImport(
+				root,
+				'strip-ansi',
+				'bun',
+				'stripANSI',
+			);
+
+			for (const importName of localNames) {
+				const callEdits = computeSimpleCallReplacementEdits(
+					root,
+					importName,
+					'stripANSI',
+				);
+				edits.push(...callEdits);
 			}
 
 			return root.commitEdits(edits);