diff --git a/codemods/invariant/index.js b/codemods/invariant/index.js index 28fe3bd..8284fdd 100644 --- a/codemods/invariant/index.js +++ b/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); }, diff --git a/codemods/iterate-iterator/index.js b/codemods/iterate-iterator/index.js index a3d476d..9810ef2 100644 --- a/codemods/iterate-iterator/index.js +++ b/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: { diff --git a/codemods/shared-ast-grep.js b/codemods/shared-ast-grep.js index 70a6498..a143894 100644 --- a/codemods/shared-ast-grep.js +++ b/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};`, ), ); } diff --git a/codemods/strip-ansi/bun/index.js b/codemods/strip-ansi/bun/index.js index 4d95df0..b884662 100644 --- a/codemods/strip-ansi/bun/index.js +++ b/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); diff --git a/codemods/strip-ansi/node/index.js b/codemods/strip-ansi/node/index.js index e7c1d10..4c6a640 100644 --- a/codemods/strip-ansi/node/index.js +++ b/codemods/strip-ansi/node/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,54 +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 { stripVTControlCharacters } from ${quoteType}node:util${quoteType};`, - ), - ); - } else { - edits.push( - imp.replace( - `const { stripVTControlCharacters } = require(${quoteType}node:util${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(`stripVTControlCharacters(${valueMatch.text()})`), - ); - } - } + const { edits, localNames } = replaceDefaultWithNamedImport( + root, + 'strip-ansi', + 'node:util', + 'stripVTControlCharacters', + ); + + for (const importName of localNames) { + const callEdits = computeSimpleCallReplacementEdits( + root, + importName, + 'stripVTControlCharacters', + ); + edits.push(...callEdits); } return root.commitEdits(edits); diff --git a/types/codemods/shared-ast-grep.d.ts b/types/codemods/shared-ast-grep.d.ts index 9fb65c0..c9fdab9 100644 --- a/types/codemods/shared-ast-grep.d.ts +++ b/types/codemods/shared-ast-grep.d.ts @@ -19,13 +19,5 @@ export function removeImport(root: SgNode, moduleName: string): { edits: Edit[]; localNames: string[]; }; -/** - * Find all named imports from a specific module in the AST. - * - * @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. - */ -export function findNamedDefaultImport(root: SgNode, moduleName: string): SgNode[]; export type SgNode = import("@ast-grep/napi").SgNode; export type Edit = import("@ast-grep/napi").Edit;