feat: add iterateComments method

src/comments.test.ts

@@ -1,9 +1,53 @@
 import ts from "typescript";
 import { describe, expect, it, vitest } from "vitest";
 
-import { forEachComment } from "./comments";
+import { forEachComment, iterateComments } from "./comments";
 import { createNodeAndSourceFile } from "./test/utils";
 
+describe("iterateComments", () => {
+	it("Should iterate all comments", () => {
+		const { node, sourceFile } = createNodeAndSourceFile(`
+			// line comment${"  "}
+			/*
+			block comment line 1${"  "}
+			block comment line 2${"  "}
+			*/
+			let value;
+		`);
+
+		const generator = iterateComments(node, sourceFile);
+		expect(typeof generator[Symbol.iterator]).toBe("function");
+		expect(generator.next()).toEqual({
+			done: false,
+			value: {
+				end: 21,
+				kind: ts.SyntaxKind.SingleLineCommentTrivia,
+				pos: 4,
+				text: "// line comment  ",
+				value: " line comment  ",
+			},
+		});
+		expect(generator.next()).toEqual({
+			done: false,
+			value: {
+				end: 85,
+				kind: ts.SyntaxKind.MultiLineCommentTrivia,
+				pos: 25,
+				text: `/*
+			block comment line 1${"  "}
+			block comment line 2${"  "}
+			*/`,
+				value: `
+			block comment line 1${"  "}
+			block comment line 2${"  "}
+			`,
+			},
+		});
+		expect(generator.next()).toEqual({ done: true, value: undefined });
+	});
+});
+
+// TODO: Move tests into `iterateComments`
 describe("forEachComment", () => {
 	it("does not call the callback when the source is a variable with no comments", () => {
 		const { node, sourceFile } = createNodeAndSourceFile("let value;");

src/comments.ts

@@ -5,6 +5,14 @@ import ts from "typescript";
 
 import { iterateTokens } from "./tokens";
 
+/**
+ * Descriptive data for a comment as yielded by {@link iterateComments}.
+ */
+export type Comment = ts.CommentRange & {
+	text: string;
+	value: string;
+};
+
 /**
  * Callback type used for {@link forEachComment}.
  * @category Callbacks
@@ -39,6 +47,28 @@ export function forEachComment(
 	callback: ForEachCommentCallback,
 	sourceFile: ts.SourceFile = node.getSourceFile(),
 ): void {
+	const fullText = sourceFile.text;
+	for (const { end, kind, pos } of iterateComments(node, sourceFile)) {
+		callback(fullText, { end, kind, pos });
+	}
+}
+
+/**
+ * Iterates over all comments owned by `node` or its children.
+ * @category Nodes - Other Utilities
+ * @example
+ * ```ts
+ * declare const node: ts.Node;
+ *
+ * for (const {pos, text} of iterateComment(node) {
+ *    console.log(`Found comment at position ${pos}: '${text}'.`);
+ * };
+ * ```
+ */
+export function* iterateComments(
+	node: ts.Node,
+	sourceFile: ts.SourceFile = node.getSourceFile(),
+): Generator<Comment> {
 	/* Visit all tokens and skip trivia.
        Comment ranges between tokens are parsed without the need of a scanner.
        forEachTokenWithWhitespace does intentionally not pay attention to the correct comment ownership of nodes as it always
@@ -53,22 +83,22 @@ export function forEachComment(
 		}
 
 		if (token.kind !== ts.SyntaxKind.JsxText) {
-			ts.forEachLeadingCommentRange(
-				fullText,
-				// skip shebang at position 0
-				token.pos === 0 ? (ts.getShebang(fullText) ?? "").length : token.pos,
-				commentCallback,
-			);
+			yield* collectComments((callback) => {
+				ts.forEachLeadingCommentRange(
+					fullText,
+					// skip shebang at position 0
+					token.pos === 0 ? (ts.getShebang(fullText) ?? "").length : token.pos,
+					callback,
+				);
+			}, fullText);
 		}
 
 		if (notJsx || canHaveTrailingTrivia(token)) {
-			ts.forEachTrailingCommentRange(fullText, token.end, commentCallback);
+			yield* collectComments((callback) => {
+				ts.forEachTrailingCommentRange(fullText, token.end, callback);
+			}, fullText);
 		}
 	}
-
-	function commentCallback(pos: number, end: number, kind: ts.CommentKind) {
-		callback(fullText, { end, kind, pos });
-	}
 }
 
 /**
@@ -105,6 +135,30 @@ function canHaveTrailingTrivia(token: ts.Node): boolean {
 	return true;
 }
 
+/**
+ * Collect comments by `ts.{forEachLeadingCommentRange,forEachTrailingCommentRange}`
+ * @internal
+ */
+function collectComments(
+	execute: (
+		callback: (pos: number, end: number, kind: ts.CommentKind) => void,
+	) => void,
+	fullText: string,
+) {
+	const comments: Comment[] = [];
+
+	execute((pos: number, end: number, kind: ts.CommentKind) => {
+		const text = fullText.slice(pos, end);
+		const value = text.slice(
+			2,
+			kind === ts.SyntaxKind.SingleLineCommentTrivia ? undefined : -2,
+		);
+		comments.push({ end, kind, pos, text, value });
+	});
+
+	return comments;
+}
+
 /**
  * Test if a node is a `JsxElement` or `JsxFragment`.
  * @internal