Merge pull request #1220 from Quramy/find_all_nodes

chore: Allow findAllNodes callback to return generic AST nodes

Author: Quramy

src/ts-ast-util/utilily-functions.test.ts

@@ -30,7 +30,7 @@ describe(isTagged, () => {
 });
 
 describe(findAllNodes, () => {
-  it('findAllNodes should return nodes which match given condition', () => {
+  it('should return nodes which match given condition', () => {
     // prettier-ignore
     const text = 'const a = `AAA`;' + '\n'
              + 'const b = `BBB`;';
@@ -39,6 +39,18 @@ describe(findAllNodes, () => {
     expect(actual.length).toBe(2);
     expect(actual.map(n => n.getText())).toEqual(['`AAA`', '`BBB`']);
   });
+
+  it('should accepts callback which returns ts.Node', () => {
+    const text = `
+      const a = fn(\`AAA\`);
+      const b = fn('BBB');
+    `;
+    const s = ts.createSourceFile('input.ts', text, ts.ScriptTarget.Latest);
+    const actual = findAllNodes(s, node => (ts.isCallExpression(node) ? node.arguments[0] : undefined));
+    expect(actual.length).toBe(2);
+    expect(ts.isNoSubstitutionTemplateLiteral(actual[0])).toBeTruthy();
+    expect(ts.isStringLiteral(actual[1])).toBeTruthy();
+  });
 });
 
 describe(isImportDeclarationWithCondition, () => {

src/ts-ast-util/utilily-functions.ts

@@ -48,18 +48,22 @@ export function findNode(sourceFile: ts.SourceFile, position: number): ts.Node |
   return find(sourceFile);
 }
 
-export function findAllNodes(sourceFile: ts.SourceFile, cond: (n: ts.Node) => boolean): ts.Node[] {
-  const result: ts.Node[] = [];
+export function findAllNodes<S extends ts.Node>(
+  sourceFile: ts.SourceFile,
+  cond: (n: ts.Node) => S | boolean | undefined,
+): S[] {
+  const result: (S | ts.Node)[] = [];
   function find(node: ts.Node) {
-    if (cond(node)) {
-      result.push(node);
+    const hit = cond(node);
+    if (hit) {
+      result.push(hit === true ? node : hit);
       return;
     } else {
       ts.forEachChild(node, find);
     }
   }
   find(sourceFile);
-  return result;
+  return result as S[];
 }
 
 export function hasTagged(node: ts.Node | undefined, condition: TagCondition) {