diff --git a/.changeset/pretty-spiders-hear.md b/.changeset/pretty-spiders-hear.md new file mode 100644 index 0000000..812f4d5 --- /dev/null +++ b/.changeset/pretty-spiders-hear.md @@ -0,0 +1,5 @@ +--- +'esrap': minor +--- + +feat: add additionalComments option for programmatic comment insertion diff --git a/README.md b/README.md index 3c7f62f..ad13df2 100644 --- a/README.md +++ b/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() }) ); ``` diff --git a/src/languages/ts/index.js b/src/languages/ts/index.js index 7bb4f7c..9543123 100644 --- a/src/languages/ts/index.js +++ b/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} */ + 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'], diff --git a/src/languages/ts/public.d.ts b/src/languages/ts/public.d.ts index 06fa86b..f76dcc2 100644 --- a/src/languages/ts/public.d.ts +++ b/src/languages/ts/public.d.ts @@ -1,2 +1,2 @@ export * from './index'; -export { Comment } from '../types'; +export { Comment, AdditionalComment } from '../types'; diff --git a/src/languages/tsx/public.d.ts b/src/languages/tsx/public.d.ts index 06fa86b..f76dcc2 100644 --- a/src/languages/tsx/public.d.ts +++ b/src/languages/tsx/public.d.ts @@ -1,2 +1,2 @@ export * from './index'; -export { Comment } from '../types'; +export { Comment, AdditionalComment } from '../types'; diff --git a/src/languages/types.d.ts b/src/languages/types.d.ts index 1229d88..3aa92d4 100644 --- a/src/languages/types.d.ts +++ b/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; }; 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'; +} diff --git a/test/additional-comments.test.js b/test/additional-comments.test.js new file mode 100644 index 0000000..fa72357 --- /dev/null +++ b/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('//'); +});