Skip to content

Commit e2ca2a2

Browse files
feat: add iterateTokens method (#756)
<!-- 👋 Hi, thanks for sending a PR to ts-api-utils! 💖. Please fill out all fields below and make sure each item is true and [x] checked. Otherwise we may not be able to review your PR. --> ## PR Checklist - [x] Addresses an existing open issue: #755 - [x] That issue was marked as [`status: accepting prs`](https://github.com/JoshuaKGoldberg/ts-api-utils/issues?q=is%3Aopen+is%3Aissue+label%3A%22status%3A+accepting+prs%22) - [x] Steps in [CONTRIBUTING.md](https://github.com/JoshuaKGoldberg/ts-api-utils/blob/main/.github/CONTRIBUTING.md) were taken ## Overview <!-- Description of what is changed and how the code change does that. --> Part of #755, this one is unlike `forEachComments`, we only need yield `token`s. We can decide what to do with the comments in a follow-up pr. --------- Co-authored-by: Josh Goldberg ✨ <git@joshuakgoldberg.com>
1 parent e547759 commit e2ca2a2

3 files changed

Lines changed: 71 additions & 26 deletions

File tree

src/comments.ts

Lines changed: 19 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33

44
import ts from "typescript";
55

6-
import { forEachToken } from "./tokens";
6+
import { iterateTokens } from "./tokens";
77

88
/**
99
* Callback type used for {@link forEachComment}.
@@ -46,32 +46,26 @@ export function forEachComment(
4646
Comment ownership is done right in this function*/
4747
const fullText = sourceFile.text;
4848
const notJsx = sourceFile.languageVariant !== ts.LanguageVariant.JSX;
49-
return forEachToken(
50-
node,
51-
(token) => {
52-
if (token.pos === token.end) {
53-
return;
54-
}
5549

56-
if (token.kind !== ts.SyntaxKind.JsxText) {
57-
ts.forEachLeadingCommentRange(
58-
fullText,
59-
// skip shebang at position 0
60-
token.pos === 0 ? (ts.getShebang(fullText) ?? "").length : token.pos,
61-
commentCallback,
62-
);
63-
}
50+
for (const token of iterateTokens(node, sourceFile)) {
51+
if (token.pos === token.end) {
52+
continue;
53+
}
54+
55+
if (token.kind !== ts.SyntaxKind.JsxText) {
56+
ts.forEachLeadingCommentRange(
57+
fullText,
58+
// skip shebang at position 0
59+
token.pos === 0 ? (ts.getShebang(fullText) ?? "").length : token.pos,
60+
commentCallback,
61+
);
62+
}
63+
64+
if (notJsx || canHaveTrailingTrivia(token)) {
65+
ts.forEachTrailingCommentRange(fullText, token.end, commentCallback);
66+
}
67+
}
6468

65-
if (notJsx || canHaveTrailingTrivia(token)) {
66-
return ts.forEachTrailingCommentRange(
67-
fullText,
68-
token.end,
69-
commentCallback,
70-
);
71-
}
72-
},
73-
sourceFile,
74-
);
7569
function commentCallback(pos: number, end: number, kind: ts.CommentKind) {
7670
callback(fullText, { end, kind, pos });
7771
}

src/tokens.test.ts

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
import ts from "typescript";
2+
import { describe, expect, it, vitest } from "vitest";
3+
4+
import { createNodeAndSourceFile } from "./test/utils";
5+
import { forEachToken, iterateTokens } from "./tokens";
6+
7+
describe("iterateTokens", () => {
8+
it("Should iterate all tokens", () => {
9+
const { node, sourceFile } = createNodeAndSourceFile("let value;");
10+
const generator = iterateTokens(node, sourceFile);
11+
expect(typeof generator[Symbol.iterator]).toBe("function");
12+
13+
const tokens = [...generator];
14+
expect(tokens.length).toBe(3);
15+
expect(tokens.every((token) => ts.isTokenKind(token.kind))).toBe(true);
16+
expect(generator.next()).toEqual({ done: true, value: undefined });
17+
});
18+
});
19+
20+
describe("forEachToken", () => {
21+
it("Should iterate all tokens", () => {
22+
const { node, sourceFile } = createNodeAndSourceFile("let value;");
23+
const callback = vitest.fn();
24+
25+
forEachToken(node, callback, sourceFile);
26+
27+
expect(callback).toBeCalledTimes(3);
28+
});
29+
});

src/tokens.ts

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,10 +36,32 @@ export function forEachToken(
3636
callback: ForEachTokenCallback,
3737
sourceFile: ts.SourceFile = node.getSourceFile(),
3838
): void {
39+
for (const token of iterateTokens(node, sourceFile)) {
40+
callback(token);
41+
}
42+
}
43+
44+
/**
45+
* Iterates over all tokens of `node`
46+
* @category Nodes - Other Utilities
47+
* @example
48+
* ```ts
49+
* declare const node: ts.Node;
50+
*
51+
* for (const token of iterateTokens(token)) {
52+
* console.log("Found token:", token.getText());
53+
* });
54+
* ```
55+
* @param node The node whose tokens should be visited
56+
*/
57+
export function* iterateTokens(
58+
node: ts.Node,
59+
sourceFile: ts.SourceFile = node.getSourceFile(),
60+
): Generator<ts.Node> {
3961
const queue = [];
4062
while (true) {
4163
if (ts.isTokenKind(node.kind)) {
42-
callback(node);
64+
yield node;
4365
} else {
4466
const children = node.getChildren(sourceFile);
4567
if (children.length === 1) {

0 commit comments

Comments
 (0)