feat: add additionalComments option for programmatic comment insertion

.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, AdditionalComment } 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) {
@@ -101,9 +101,42 @@ export default (options = {}) => {
 	const quote_char = options.quotes === 'double' ? '"' : "'";
 
 	const comments = options.comments ?? [];
+	/** @type {WeakMap<TSESTree.Node, AdditionalComment[]>} */
+	const additionalComments = options.additionalComments ?? new WeakMap();
 
 	let comment_index = 0;
 
+	/**
+	 * Write additional comments for a node
+	 * @param {Context} context
+	 * @param {TSESTree.Node} node
+	 * @param {('leading' | 'trailing')} position
+	 */
+	function write_additional_comments(context, node, position) {
+		const nodeComments = additionalComments.get(node);
+		if (!nodeComments) return;
+
+		const relevantComments = nodeComments.filter((comment) => comment.position === position);
+
+		for (let i = 0; i < relevantComments.length; i += 1) {
+			const comment = relevantComments[i];
+
+			if (position === 'trailing' && i === 0) {
+				context.write(' ');
+			}
+
+			write_comment(comment, context);
+
+			if (position === 'leading') {
+				if (comment.type === 'Line') {
+					context.newline();
+				} else {
+					context.write(' ');
+				}
+			}
+		}
+	}
+
 	/**
 	 * Set `comment_index` to be the first comment after `start`.
 	 * Most of the time this is already correct, but if nodes
@@ -586,11 +619,15 @@ export default (options = {}) => {
 
 	return {
 		_(node, context, visit) {
+			write_additional_comments(context, node, 'leading');
+
 			if (node.loc) {
 				flush_comments_until(context, null, node.loc.start, true);
 			}
 
 			visit(node);
+
+			write_additional_comments(context, node, 'trailing');
 		},
 
 		ArrayExpression: shared['ArrayExpression|ArrayPattern'],

src/languages/ts/public.d.ts

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

src/languages/tsx/public.d.ts

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

src/languages/types.d.ts

@@ -1,6 +1,9 @@
+import { TSESTree } from '@typescript-eslint/types';
+
 export type TSOptions = {
 	quotes?: 'double' | 'single';
 	comments?: Comment[];
+	additionalComments?: WeakMap<TSESTree.Node, AdditionalComment[]>;
 };
 
 interface Position {
@@ -10,13 +13,20 @@ 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;
 	};
 }
+
+export interface AdditionalComment extends BaseComment {
+	position: 'leading' | 'trailing';
+}

test/additional-comments.test.js

@@ -0,0 +1,105 @@
+// @ts-check
+/** @import { TSESTree } from '@typescript-eslint/types' */
+/** @import { AdditionalComment } from '../src/languages/types.js' */
+import { expect, test } from 'vitest';
+import { print } from '../src/index.js';
+import { load } from './common.js';
+import ts from '../src/languages/ts/index.js';
+
+/**
+ * Helper to create additional comments and print code
+ * @param {TSESTree.Program} ast - Parsed AST
+ * @param {TSESTree.Node} node - AST node to attach comments to
+ * @param {AdditionalComment[]} comments - Comments to attach
+ * @returns {string} Generated code
+ */
+function print_with_comments(ast, node, comments) {
+	const additionalComments = new WeakMap();
+	additionalComments.set(node, comments);
+
+	const output = print(ast, ts({ additionalComments }));
+	return output.code;
+}
+
+/**
+ * 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 } = load(input);
+	const returnStatement = get_return_statement(ast);
+	expect(returnStatement.type).toBe('ReturnStatement');
+
+	/** @type {AdditionalComment[]} */
+	const comments = [
+		{
+			type: 'Line',
+			value: ' This is a leading comment',
+			position: 'leading'
+		},
+		{
+			type: 'Block',
+			value: ' This is a trailing comment ',
+			position: 'trailing'
+		}
+	];
+
+	const code = print_with_comments(ast, returnStatement, comments);
+
+	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 } = load(input);
+	const returnStatement = get_return_statement(ast);
+
+	/** @type {AdditionalComment[]} */
+	const comments = [
+		{
+			type: 'Line',
+			value: ' Leading only',
+			position: 'leading'
+		}
+	];
+
+	const code = print_with_comments(ast, returnStatement, comments);
+
+	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 } = load(input);
+	const returnStatement = get_return_statement(ast);
+
+	/** @type {AdditionalComment[]} */
+	const comments = [
+		{
+			type: 'Block',
+			value: ' Trailing only ',
+			position: 'trailing'
+		}
+	];
+
+	const code = print_with_comments(ast, returnStatement, comments);
+
+	expect(code).toContain('/* Trailing only */');
+	expect(code).not.toContain('//');
+});