feat: add iterateTokens method

src/comments.ts

@@ -3,7 +3,7 @@
 
 import ts from "typescript";
 
-import { forEachToken } from "./tokens";
+import { iterateTokens } from "./tokens";
 
 /**
  * Callback type used for {@link forEachComment}.
@@ -46,32 +46,26 @@ export function forEachComment(
        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) {
-				ts.forEachLeadingCommentRange(
-					fullText,
-					// skip shebang at position 0
-					token.pos === 0 ? (ts.getShebang(fullText) ?? "").length : token.pos,
-					commentCallback,
-				);
-			}
+	for (const token of iterateTokens(node, sourceFile)) {
+		if (token.pos === token.end) {
+			continue;
+		}
+
+		if (token.kind !== ts.SyntaxKind.JsxText) {
+			ts.forEachLeadingCommentRange(
+				fullText,
+				// skip shebang at position 0
+				token.pos === 0 ? (ts.getShebang(fullText) ?? "").length : token.pos,
+				commentCallback,
+			);
+		}
+
+		if (notJsx || canHaveTrailingTrivia(token)) {
+			ts.forEachTrailingCommentRange(fullText, token.end, commentCallback);
+		}
+	}
 
-			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 });
 	}

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) {