feat: getLeadingComments/getTrailingComments

.changeset/pretty-spiders-hear.md

@@ -0,0 +1,5 @@
+---
+'esrap': minor
+---
+
+feat: add additionalComments option for programmatic comment insertion

README.md

@@ -55,7 +55,12 @@ const { code, map } = print(
     quotes: 'single',
 
     // an array of `{ type: 'Line' | 'Block', value: string, loc: { start, end } }` objects
-    comments: []
+    comments: [],
+
+    // a WeakMap of AST nodes to additional comments to insert at specific nodes
+    // useful for programmatically adding comments during code transformation,
+    // especially for nodes that were added programmatically
+    additionalComments: new WeakMap()
   })
 );
 ```

src/languages/ts/index.js

@@ -1,6 +1,6 @@
 /** @import { TSESTree } from '@typescript-eslint/types' */
 /** @import { Visitors } from '../../types.js' */
-/** @import { TSOptions, Comment } from '../types.js' */
+/** @import { TSOptions, BaseComment } from '../types.js' */
 import { Context } from 'esrap';
 
 /** @typedef {TSESTree.Node} Node */
@@ -74,7 +74,7 @@ const OPERATOR_PRECEDENCE = {
 };
 
 /**
- * @param {Comment} comment
+ * @param {BaseComment} comment
  * @param {Context} context
  */
 function write_comment(comment, context) {
@@ -90,6 +90,7 @@ function write_comment(comment, context) {
 		}
 
 		context.write('*/');
+		if (lines.length > 1) context.newline();
 	}
 }
 
@@ -104,6 +105,36 @@ export default (options = {}) => {
 
 	let comment_index = 0;
 
+	/**
+	 * Write additional comments for a node
+	 * @param {Context} context
+	 * @param {BaseComment[] | undefined} comments
+	 * @param {('leading' | 'trailing')} position
+	 */
+	function write_additional_comments(context, comments, position) {
+		if (!comments) {
+			return;
+		}
+
+		for (let i = 0; i < comments.length; i += 1) {
+			const comment = comments[i];
+
+			if (position === 'trailing' && i === 0) {
+				context.write(' ');
+			}
+
+			write_comment(comment, context);
+
+			if (position === 'leading') {
+				if (comment.type === 'Line') {
+					context.newline();
+				} else if (comment.type === 'Block' && !comment.value.includes('\n')) {
+					context.write(' ');
+				}
+			}
+		}
+	}
+
 	/**
 	 * Set `comment_index` to be the first comment after `start`.
 	 * Most of the time this is already correct, but if nodes
@@ -775,11 +806,15 @@ export default (options = {}) => {
 
 	return {
 		_(node, context, visit) {
+			write_additional_comments(context, options.getLeadingComments?.(node), 'leading');
+
 			if (node.loc) {
 				flush_comments_until(context, null, node.loc.start, true);
 			}
 
 			visit(node);
+
+			write_additional_comments(context, options.getTrailingComments?.(node), 'trailing');
 		},
 
 		AccessorProperty:

src/languages/ts/public.d.ts

@@ -1,2 +1,2 @@
 export * from './index';
-export { Comment } from '../types';
+export { BaseComment, Comment } from '../types';

src/languages/tsx/public.d.ts

@@ -1,2 +1,2 @@
 export * from './index';
-export { Comment } from '../types';
+export { BaseComment, Comment } from '../types';

src/languages/types.d.ts

@@ -1,6 +1,10 @@
+import { TSESTree } from '@typescript-eslint/types';
+
 export type TSOptions = {
 	quotes?: 'double' | 'single';
 	comments?: Comment[];
+	getLeadingComments?: (node: TSESTree.Node) => BaseComment[] | undefined;
+	getTrailingComments?: (node: TSESTree.Node) => BaseComment[] | undefined;
 };
 
 interface Position {
@@ -10,11 +14,14 @@ interface Position {
 
 // this exists in TSESTree but because of the inanity around enums
 // it's easier to do this ourselves
-export interface Comment {
+export interface BaseComment {
 	type: 'Line' | 'Block';
 	value: string;
 	start?: number;
 	end?: number;
+}
+
+export interface Comment extends BaseComment {
 	loc: {
 		start: Position;
 		end: Position;

test/additional-comments.test.js

@@ -0,0 +1,122 @@
+// @ts-check
+/** @import { TSESTree } from '@typescript-eslint/types' */
+/** @import { BaseComment } from '../src/languages/types.js' */
+import { expect, test } from 'vitest';
+import { print } from '../src/index.js';
+import { acornParse } from './common.js';
+import ts from '../src/languages/ts/index.js';
+
+/**
+ * @param {string} value
+ * @returns {BaseComment}
+ */
+function line(value) {
+	return { type: 'Line', value };
+}
+
+/**
+ * @param {string} value
+ * @returns {BaseComment}
+ */
+function block(value) {
+	return { type: 'Block', value };
+}
+
+/**
+ * Helper to get return statement from a simple function
+ * @param {TSESTree.Program} ast - Parsed AST
+ * @returns {TSESTree.Node} The return statement
+ */
+function get_return_statement(ast) {
+	const functionDecl = ast.body[0];
+	// @ts-expect-error accessing function body
+	const statements = functionDecl.body.body;
+	// Find the return statement (could be first or second depending on function structure)
+	return statements.find(/** @param {any} stmt */ (stmt) => stmt.type === 'ReturnStatement');
+}
+
+test('additional comments are inserted correctly', () => {
+	const input = `function example() {
+	const x = 1;
+	return x;
+}`;
+
+	const { ast } = acornParse(input);
+	const returnStatement = get_return_statement(ast);
+	expect(returnStatement.type).toBe('ReturnStatement');
+
+	const { code } = print(
+		ast,
+		ts({
+			getLeadingComments: (n) =>
+				n === returnStatement ? [line(' This is a leading comment')] : undefined,
+			getTrailingComments: (n) =>
+				n === returnStatement ? [block(' This is a trailing comment ')] : undefined
+		})
+	);
+
+	expect(code).toContain('// This is a leading comment');
+	expect(code).toContain('/* This is a trailing comment */');
+});
+
+test('only leading comments are inserted when specified', () => {
+	const input = `function test() { return 42; }`;
+	const { ast } = acornParse(input);
+	const returnStatement = get_return_statement(ast);
+
+	const { code } = print(
+		ast,
+		ts({
+			getLeadingComments: (n) => (n === returnStatement ? [line(' Leading only ')] : undefined)
+		})
+	);
+
+	expect(code).toContain('// Leading only');
+	expect(code).not.toContain('trailing');
+});
+
+test('only trailing comments are inserted when specified', () => {
+	const input = `function test() { return 42; }`;
+	const { ast } = acornParse(input);
+	const returnStatement = get_return_statement(ast);
+
+	const { code } = print(
+		ast,
+		ts({
+			getTrailingComments: (n) => (n === returnStatement ? [block(' Trailing only ')] : undefined)
+		})
+	);
+
+	expect(code).toContain('/* Trailing only */');
+	expect(code).not.toContain('//');
+});
+
+test('additional comments multi-line comments have new line', () => {
+	const input = `function example() {
+	const x = 1;
+	return x;
+}`;
+
+	const { ast } = acornParse(input);
+	const returnStatement = get_return_statement(ast);
+	expect(returnStatement.type).toBe('ReturnStatement');
+
+	const { code } = print(
+		ast,
+		ts({
+			getLeadingComments: (n) =>
+				n === returnStatement ? [block('*\n * This is a leading comment\n ')] : undefined
+		})
+	);
+
+	expect(code).toMatchInlineSnapshot(`
+		"function example() {
+			const x = 1;
+
+			/**
+			 * This is a leading comment
+			 */
+			return x;
+		}"
+	`);
+});