feat(utils): add shebang

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\'');
+		});
+	});
+});

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<string, string>, 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;
+		}
+	}
+};