feat: add iterateComments method

src/comments.test.ts

@@ -1,9 +1,33 @@
 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(`
+			// hello world
+			let value;
+		`);
+
+		const generator = iterateComments(node, sourceFile);
+		expect(typeof generator[Symbol.iterator]).toBe("function");
+
+		expect([...generator]).toEqual([
+			{
+				end: 18,
+				fullText: sourceFile.getFullText(),
+				kind: ts.SyntaxKind.SingleLineCommentTrivia,
+				pos: 4,
+				text: "// hello world",
+			},
+		]);
+		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

@@ -3,7 +3,11 @@
 
 import ts from "typescript";
 
-import { forEachToken } from "./tokens";
+import { iterateTokens } from "./tokens";
+
+export type Comment = ts.CommentRange & {
+	fullText: string;
+};
 
 /**
  * Callback type used for {@link forEachComment}.
@@ -39,41 +43,59 @@ export function forEachComment(
 	callback: ForEachCommentCallback,
 	sourceFile: ts.SourceFile = node.getSourceFile(),
 ): void {
+	for (const { end, fullText, 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
        scans all trivia before each token, which could include trailing comments of the previous token.
        Comment ownership is done right in this function*/
 	const fullText = sourceFile.text;
 	const notJsx = sourceFile.languageVariant !== ts.LanguageVariant.JSX;
-	return forEachToken(
-		node,
-		(token) => {
-			if (token.pos === token.end) {
-				return;
-			}
 
-			if (token.kind !== ts.SyntaxKind.JsxText) {
+	for (const token of iterateTokens(node, sourceFile)) {
+		if (token.pos === token.end) {
+			continue;
+		}
+
+		if (token.kind !== ts.SyntaxKind.JsxText) {
+			yield* collectComments((callback) => {
 				ts.forEachLeadingCommentRange(
 					fullText,
 					// skip shebang at position 0
 					token.pos === 0 ? (ts.getShebang(fullText) ?? "").length : token.pos,
-					commentCallback,
+					callback,
 				);
-			}
+			}, fullText);
+		}
 
-			if (notJsx || canHaveTrailingTrivia(token)) {
-				return ts.forEachTrailingCommentRange(
-					fullText,
-					token.end,
-					commentCallback,
-				);
-			}
-		},
-		sourceFile,
-	);
-	function commentCallback(pos: number, end: number, kind: ts.CommentKind) {
-		callback(fullText, { end, kind, pos });
+		if (notJsx || canHaveTrailingTrivia(token)) {
+			yield* collectComments((callback) => {
+				ts.forEachTrailingCommentRange(fullText, token.end, callback);
+			}, fullText);
+		}
 	}
 }
 
@@ -111,6 +133,25 @@ 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) => {
+		comments.push({ end, fullText, kind, pos });
+	});
+
+	return comments;
+}
+
 /**
  * Test if a node is a `JsxElement` or `JsxFragment`.
  * @internal

src/tokens.test.ts

@@ -0,0 +1,29 @@
+import ts from "typescript";
+import { describe, expect, it, vitest } from "vitest";
+
+import { createNodeAndSourceFile } from "./test/utils";
+import { forEachToken, iterateTokens } from "./tokens";
+
+describe("iterateTokens", () => {
+	it("Should iterate all tokens", () => {
+		const { node, sourceFile } = createNodeAndSourceFile("let value;");
+		const generator = iterateTokens(node, sourceFile);
+		expect(typeof generator[Symbol.iterator]).toBe("function");
+
+		const tokens = [...generator];
+		expect(tokens.length).toBe(3);
+		expect(tokens.every((token) => ts.isTokenKind(token.kind))).toBe(true);
+		expect(generator.next()).toEqual({ done: true, value: undefined });
+	});
+});
+
+describe("forEachToken", () => {
+	it("Should iterate all tokens", () => {
+		const { node, sourceFile } = createNodeAndSourceFile("let value;");
+		const callback = vitest.fn();
+
+		forEachToken(node, callback, sourceFile);
+
+		expect(callback).toBeCalledTimes(3);
+	});
+});

src/tokens.ts

@@ -36,10 +36,32 @@ export function forEachToken(
 	callback: ForEachTokenCallback,
 	sourceFile: ts.SourceFile = node.getSourceFile(),
 ): void {
+	for (const token of iterateTokens(node, sourceFile)) {
+		callback(token);
+	}
+}
+
+/**
+ * Iterates over all tokens of `node`
+ * @category Nodes - Other Utilities
+ * @example
+ * ```ts
+ * declare const node: ts.Node;
+ *
+ * for (const token of iterateTokens(token)) {
+ * 	console.log("Found token:", token.getText());
+ * });
+ * ```
+ * @param node The node whose tokens should be visited
+ */
+export function* iterateTokens(
+	node: ts.Node,
+	sourceFile: ts.SourceFile = node.getSourceFile(),
+): Generator<ts.Node> {
 	const queue = [];
 	while (true) {
 		if (ts.isTokenKind(node.kind)) {
-			callback(node);
+			yield node;
 		} else {
 			const children = node.getChildren(sourceFile);
 			if (children.length === 1) {