feat: #279 - Add StatementedNode.getDescendantStatements().

Not so efficient at the moment, but good enough for now.
dsherret · Mar 18, 2018 · 3b8b093 · 3b8b093
1 parent f8676f0
commit 3b8b093
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 0 deletions.
diff --git a/src/compiler/statement/StatementedNode.ts b/src/compiler/statement/StatementedNode.ts
@@ -27,6 +27,10 @@ export interface StatementedNode {
      * Gets the node's statements.
      */
     getStatements(): Statement[];
+    /**
+     * Gets the node's descendant statements.
+     */
+    getDescendantStatements(): Statement[];
     /**
      * Gets the first statement that matches the provided condition or returns undefined if it doesn't exist.
      * @param findFunction - Function to find the statement by.
@@ -442,6 +446,30 @@ export function StatementedNode<T extends Constructor<StatementedNodeExtensionTy
             return this.getCompilerStatements().map(s => this.getNodeFromCompilerNode<Statement>(s));
         }
 
+        getDescendantStatements(): Statement[] {
+            // todo: Make this more efficient.
+            type NodeWithStatements = ts.Node & { statements: ts.NodeArray<ts.Statement>; };
+            return ArrayUtils.from(getDescendantStatements(this));
+
+            function* getDescendantStatements(thisNode: Node & StatementedNode) {
+                for (const childStatement of thisNode.getStatements()) {
+                    yield childStatement;
+
+                    for (const descendant of childStatement.getCompilerDescendantsIterator()) {
+                        if ((descendant as NodeWithStatements).statements != null) {
+                            for (const statement of (descendant as NodeWithStatements).statements)
+                                yield thisNode.getNodeFromCompilerNode<Statement>(statement);
+                        }
+                        else if (descendant.kind === SyntaxKind.ArrowFunction) {
+                            const arrowFunction = (descendant as ts.ArrowFunction);
+                            if (arrowFunction.body.kind !== SyntaxKind.Block)
+                                yield thisNode.getNodeFromCompilerNode<Statement>(arrowFunction.body);
+                        }
+                    }
+                }
+            }
+        }
+
         getStatement(findFunction: (statement: Node) => boolean) {
             return ArrayUtils.find(this.getStatements(), findFunction);
         }

diff --git a/src/tests/compiler/statement/statementedNodeTests.ts b/src/tests/compiler/statement/statementedNodeTests.ts
@@ -48,6 +48,39 @@ describe(nameof(StatementedNode), () => {
         });
     });
 
+    describe(nameof<StatementedNode>(s => s.getStatements), () => {
+        function doTest(text: string, expectedStatements: string[]) {
+            const {sourceFile} = getInfoFromText(text);
+            expect(sourceFile.getDescendantStatements().map(s => s.getText())).to.deep.equal(expectedStatements);
+        }
+
+        it("should get the descendant statements", () => {
+            const expected = [
+                `const a = () => {\n    const b = "";\n};`,
+                `const b = "";`,
+                `const c = 5;`,
+                `function d() {\n    function e() {\n        const f = "";\n    }\n}`,
+                `function e() {\n        const f = "";\n    }`,
+                `const f = "";`,
+                `class MyClass {\n    prop = () => console.log("here");\n}`,
+                `console.log("here")`
+            ];
+            doTest(`const a = () => {
+    const b = "";
+};
+const c = 5;
+function d() {
+    function e() {
+        const f = "";
+    }
+}
+class MyClass {
+    prop = () => console.log("here");
+}
+`, expected);
+        });
+    });
+
     describe(nameof<StatementedNode>(s => s.getStatement), () => {
         it("should get the statement when it exists", () => {
             const {sourceFile} = getInfoFromText("var t; class T {}");