function list and of type node fixes

romgerman · Apr 21, 2024 · 2adabab · 2adabab
1 parent aa8b854
commit 2adabab
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 7 deletions.
diff --git a/src/blueprint/nodes/aggregation/function-list.ts b/src/blueprint/nodes/aggregation/function-list.ts
@@ -17,7 +17,7 @@ export class FunctionListNode extends BlueprintNode {
   }
 
   private getFunctionList(tsFileList: Array<ts.SourceFile | ts.FunctionDeclaration>): ts.FunctionDeclaration[] {
-    if (isArrayOfType(tsFileList, ts.isClassDeclaration)) {
+    if (isArrayOfType(tsFileList, ts.isFunctionDeclaration)) {
       return tsFileList as ts.FunctionDeclaration[];
     }
 

diff --git a/src/blueprint/nodes/filtering/of-type-predicate.ts b/src/blueprint/nodes/filtering/of-type-predicate.ts
@@ -1,6 +1,7 @@
 import * as ts from "typescript";
 import { NodeTypes } from "../../../shared/node-types";
 import { PredicateNode } from "./filter-by-node";
+import { isArrayOfType } from "../../helpers";
 
 export enum TsNodeType {
   ClassDeclaration = "class-decl",
@@ -21,18 +22,40 @@ export class OfTypePredicateNode extends PredicateNode<{ type: string }> {
         throw new Error("Expected Node[] at input 0");
       }
 
-      if (this.state.type === TsNodeType.ClassDeclaration) {
-        return array.filter((n) => ts.isClassDeclaration(n));
-      } else if (this.state.type === TsNodeType.FunctionDeclaration) {
-        return array.filter((n) => ts.isFunctionDeclaration(n));
-      } else if (this.state.type === TsNodeType.MethodDeclaration) {
-        return array.filter((n) => ts.isMethodDeclaration(n));
+      if (isArrayOfType(array, ts.isSourceFile)) {
+        if (this.state.type === TsNodeType.ClassDeclaration) {
+          return this.getNodesOfTypeFromSourceFiles(array, ts.isClassDeclaration);
+        } else if (this.state.type === TsNodeType.FunctionDeclaration) {
+          return this.getNodesOfTypeFromSourceFiles(array, ts.isFunctionDeclaration);
+        } else if (this.state.type === TsNodeType.MethodDeclaration) {
+          return this.getNodesOfTypeFromSourceFiles(array, ts.isMethodDeclaration);
+        }
+      } else {
+        if (this.state.type === TsNodeType.ClassDeclaration) {
+          return array.filter((n) => ts.isClassDeclaration(n));
+        } else if (this.state.type === TsNodeType.FunctionDeclaration) {
+          return array.filter((n) => ts.isFunctionDeclaration(n));
+        } else if (this.state.type === TsNodeType.MethodDeclaration) {
+          return array.filter((n) => ts.isMethodDeclaration(n));
+        }
       }
 
       return [];
     };
   }
 
+  private getNodesOfTypeFromSourceFiles(sourceFiles: ts.Node[], predicate: Function) {
+    const result: ts.Node[] = [];
+    for (let sourceFile of sourceFiles) {
+      ts.forEachChild(sourceFile, (node) => {
+        if (predicate(node)) {
+          result.push(node);
+        }
+      });
+    }
+    return result;
+  }
+
   async getViewData(): Promise<any> {
     return [];
   }