diff --git a/utils/src/ast-grep/shebang.test.ts b/utils/src/ast-grep/shebang.test.ts new file mode 100644 index 00000000..b68ef9e0 --- /dev/null +++ b/utils/src/ast-grep/shebang.test.ts @@ -0,0 +1,175 @@ +import assert from "node:assert/strict"; +import { describe, it } from "node:test"; +import astGrep from "@ast-grep/napi"; +import dedent from "dedent"; +import type { Edit } from "@ast-grep/napi"; +import { getShebang, replaceNodeJsArgs, } from './shebang.ts'; + +describe("shebang", () => { + describe("getShebang", () => { + it("should get the shebang line", () => { + const code = dedent` + #!/usr/bin/env node + console.log("Hello, world!"); + `; + const ast = astGrep.parse(astGrep.Lang.JavaScript, code); + + const shebang = getShebang(ast); + + assert.ok(shebang); + assert.equal(shebang.text(), "#!/usr/bin/env node"); + }); + + it("should take the last shebang line if multiple exist on top of the code", () => { + const code = dedent` + #!/usr/bin/env node 1 + #!/usr/bin/env node 2 + console.log("Hello, world!"); + `; + const ast = astGrep.parse(astGrep.Lang.JavaScript, code); + + const shebang = getShebang(ast); + + assert.strictEqual(shebang?.text(), "#!/usr/bin/env node 2"); + }); + + it("should return null if no shebang line", () => { + const code = dedent` + console.log("Hello, world!"); + `; + + const ast = astGrep.parse(astGrep.Lang.JavaScript, code); + + const shebang = getShebang(ast); + assert.strictEqual(shebang, null); + }); + + it("shouldn't catch shebangs in comments", () => { + const code = dedent` + // #!/usr/bin/env node + console.log("Hello, world!"); + `; + const ast = astGrep.parse(astGrep.Lang.JavaScript, code); + + const shebang = getShebang(ast); + + assert.strictEqual(shebang, null); + }); + + it("shouldn't catch shebang in middle of code", () => { + const code = dedent` + console.log("Hello, world!"); + #!/usr/bin/env node + `; + const ast = astGrep.parse(astGrep.Lang.JavaScript, code); + + const shebang = getShebang(ast); + + assert.strictEqual(shebang, null); + }); + }); + + describe("replaceNodeJsArgs", () => { + it("should replace multiple different arguments in shebang with overlapping names", () => { + const code = dedent` + #!/usr/bin/env node --foo --foobar --bar + console.log("Hello, world!"); + `; + const ast = astGrep.parse(astGrep.Lang.JavaScript, code); + const edits: Edit[] = []; + + replaceNodeJsArgs(ast, { '--foo': '--baz', '--bar': '--qux' }, edits); + + assert.strictEqual(edits.length, 2); + assert.strictEqual(edits[0].insertedText, '#!/usr/bin/env node --baz --foobar --bar'); + assert.strictEqual(edits[1].insertedText, '#!/usr/bin/env node --baz --foobar --qux'); + }); + + it("should not replace arguments that are substrings of other args", () => { + const code = dedent` + #!/usr/bin/env node --foo --foo-bar --bar + console.log("Hello, world!"); + `; + const ast = astGrep.parse(astGrep.Lang.JavaScript, code); + const edits: Edit[] = []; + + replaceNodeJsArgs(ast, { '--foo': '--baz', '--bar': '--qux' }, edits); + + assert.strictEqual(edits.length, 2); + assert.strictEqual(edits[0].insertedText, '#!/usr/bin/env node --baz --foo-bar --bar'); + assert.strictEqual(edits[1].insertedText, '#!/usr/bin/env node --baz --foo-bar --qux'); + }); + + it("should handle shebang with multiple spaces between args", () => { + const code = dedent` + #!/usr/bin/env node --foo --bar + console.log("Hello, world!"); + `; + const ast = astGrep.parse(astGrep.Lang.JavaScript, code); + const edits: Edit[] = []; + + replaceNodeJsArgs(ast, { '--foo': '--baz', '--bar': '--qux' }, edits); + + assert.strictEqual(edits.length, 2); + assert.strictEqual(edits[0].insertedText, '#!/usr/bin/env node --baz --bar'); + assert.strictEqual(edits[1].insertedText, '#!/usr/bin/env node --baz --qux'); + }); + + it("should not replace if argument is at the start of the shebang", () => { + const code = dedent` + #!/usr/bin/env --foo node --bar + console.log("Hello, world!"); + `; + const ast = astGrep.parse(astGrep.Lang.JavaScript, code); + const edits: Edit[] = []; + + replaceNodeJsArgs(ast, { '--foo': '--baz' }, edits); + + // Should not replace because node must be present + assert.strictEqual(edits.length, 0); + }); + + it("should replace argument with special characters", () => { + const code = dedent` + #!/usr/bin/env node --foo-bar --bar_foo + console.log("Hello, world!"); + `; + const ast = astGrep.parse(astGrep.Lang.JavaScript, code); + const edits: Edit[] = []; + + replaceNodeJsArgs(ast, { '--foo-bar': '--baz-bar', '--bar_foo': '--qux_foo' }, edits); + + assert.strictEqual(edits.length, 2); + assert.strictEqual(edits[0].insertedText, '#!/usr/bin/env node --baz-bar --bar_foo'); + assert.strictEqual(edits[1].insertedText, '#!/usr/bin/env node --baz-bar --qux_foo'); + }); + + it("should not replace anything if argsToValues is empty", () => { + const code = dedent` + #!/usr/bin/env node --foo --bar + console.log("Hello, world!"); + `; + const ast = astGrep.parse(astGrep.Lang.JavaScript, code); + const edits: Edit[] = []; + + replaceNodeJsArgs(ast, {}, edits); + + assert.strictEqual(edits.length, 0); + }); + + it("should handle shebang with quoted arguments", () => { + const code = dedent` + #!/usr/bin/env node "--foo" '--bar' + console.log("Hello, world!"); + `; + const ast = astGrep.parse(astGrep.Lang.JavaScript, code); + const edits: Edit[] = []; + + replaceNodeJsArgs(ast, { '"--foo"': '"--baz"', "'--bar'": "'--qux'" }, edits); + + assert.strictEqual(edits.length, 2); + assert.strictEqual(edits[0].insertedText, '#!/usr/bin/env node "--baz" \'--bar\''); + assert.strictEqual(edits[1].insertedText, '#!/usr/bin/env node "--baz" \'--qux\''); + }); + }); +}); diff --git a/utils/src/ast-grep/shebang.ts b/utils/src/ast-grep/shebang.ts new file mode 100644 index 00000000..6c880c69 --- /dev/null +++ b/utils/src/ast-grep/shebang.ts @@ -0,0 +1,65 @@ +import type { SgRoot, Edit } from "@codemod.com/jssg-types/main"; + +/** + * Get the shebang line from the root. + * @param root The root node to search. + * @returns The shebang line if found, otherwise null. + */ +export const getShebang = (root: SgRoot) => + root + .root() + .find({ + rule: { + kind: "hash_bang_line", + regex: "\\bnode(\\.exe)?\\b", + not: { + // tree-sitter wrap hash bang in Error node + // when it's not in the top of program node + inside: { + kind: "ERROR" + } + } + } + }) + +/** + * Replace Node.js arguments in the shebang line. + * @param root The root node to search. + * @param argsToValues The mapping of argument names to their new values. + * @param edits The list of edits to apply. + * @returns The updated shebang line if any replacements were made, otherwise null. + */ +export const replaceNodeJsArgs = (root: SgRoot, argsToValues: Record, edits: Edit[]) => { + const shebang = getShebang(root); + + if (!shebang) return; + + const text = shebang.text(); + // Find the "node" argument in the shebang + const nodeMatch = text.match(/\bnode(\.exe)?\b/); + + if (!nodeMatch) return; + + // We only touch to something after node because before it's env thing + const nodeIdx = nodeMatch.index! + nodeMatch[0].length; + const beforeNode = text.slice(0, nodeIdx); + let afterNode = text.slice(nodeIdx); + + for (const argC of Object.keys(argsToValues)) { + // Escape special regex characters in arg + const esc = argC.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"); + const regex = new RegExp(`(\\s+)(["']?)${esc}(["']?)(?=\\s|$)`, 'g'); + + // handling quote and whitespaces + const newAfterNode = afterNode.replace(regex, (_unused, ws, q1, q2) => { + const replacement = argsToValues[argC]; + + return `${ws}${q1}${replacement}${q2}`; + }); + + if (newAfterNode !== afterNode) { + edits.push(shebang.replace(beforeNode + newAfterNode)); + afterNode = newAfterNode; + } + } +};