feat: #506 - Get variable statement by a declaration's name.

dsherret · Dec 2, 2018 · 6b94548 · 6b94548
1 parent 8a8b555
commit 6b94548
Show file tree

Hide file tree

Showing 2 changed files with 68 additions and 14 deletions.
diff --git a/src/compiler/ast/statement/StatementedNode.ts b/src/compiler/ast/statement/StatementedNode.ts
@@ -371,16 +371,30 @@ export interface StatementedNode {
      * Gets the direct variable statement children.
      */
     getVariableStatements(): VariableStatement[];
+    /**
+     * Gets a variable statement.
+     * @param name - Name of one of the variable statement's declarations.
+     */
+    getVariableStatement(name: string): VariableStatement | undefined;
     /**
      * Gets a variable statement.
      * @param findFunction - Function to use to find the variable statement.
      */
     getVariableStatement(findFunction: (declaration: VariableStatement) => boolean): VariableStatement | undefined;
+    /** @internal */
+    getVariableStatement(nameOrFindFunction: string | ((declaration: VariableStatement) => boolean)): VariableStatement | undefined;
+    /**
+     * Gets a variable statement or throws if it doesn't exist.
+     * @param name - Name of one of the variable statement's declarations.
+     */
+    getVariableStatementOrThrow(name: string): VariableStatement;
     /**
      * Gets a variable statement or throws if it doesn't exist.
      * @param findFunction - Function to use to find the variable statement.
      */
     getVariableStatementOrThrow(findFunction: (declaration: VariableStatement) => boolean): VariableStatement;
+    /** @internal */
+    getVariableStatementOrThrow(nameOrFindFunction: string | ((declaration: VariableStatement) => boolean)): VariableStatement;
     /**
      * Gets all the variable declarations within all the variable declarations of the direct variable statement children.
      */
@@ -791,12 +805,18 @@ export function StatementedNode<T extends Constructor<StatementedNodeExtensionTy
             return childSyntaxList.getChildrenOfKind(SyntaxKind.VariableStatement);
         }
 
-        getVariableStatement(findFunction: (declaration: VariableStatement) => boolean): VariableStatement | undefined {
-            return ArrayUtils.find(this.getVariableStatements(), findFunction);
+        getVariableStatement(nameOrFindFunction: string | ((statement: VariableStatement) => boolean)): VariableStatement | undefined {
+            return ArrayUtils.find(this.getVariableStatements(), getFindFunction());
+
+            function getFindFunction() {
+                if (typeof nameOrFindFunction === "string")
+                    return (statement: VariableStatement) => statement.getDeclarations().some(d => d.getName() === nameOrFindFunction);
+                return nameOrFindFunction;
+            }
         }
 
-        getVariableStatementOrThrow(findFunction: (declaration: VariableStatement) => boolean): VariableStatement {
-            return errors.throwIfNullOrUndefined(this.getVariableStatement(findFunction), "Expected to find a variable statement that matched the provided condition.");
+        getVariableStatementOrThrow(nameOrFindFunction: string | ((statement: VariableStatement) => boolean)): VariableStatement {
+            return errors.throwIfNullOrUndefined(this.getVariableStatement(nameOrFindFunction), "Expected to find a variable statement that matched the provided condition.");
         }
 
         addVariableStatement(structure: VariableStatementStructure) {

diff --git a/src/tests/compiler/statement/statementedNode/variableTests.ts b/src/tests/compiler/statement/statementedNode/variableTests.ts
@@ -138,24 +138,58 @@ describe(nameof(StatementedNode), () => {
     });
 
     describe(nameof<StatementedNode>(n => n.getVariableStatement), () => {
-        it("should get a variable statement when something matches", () => {
-            const statement = variablesSourceFile.getVariableStatement(s => s.getDeclarations().length === 2)!;
-            expect(statement.getDeclarations()[0].getName()).to.equal("Identifier2");
+        function doTest(nameOrFindFunction: string | ((declaration: VariableStatement) => boolean), expectedFirstDeclarationName: string | undefined) {
+            const statement = variablesSourceFile.getVariableStatement(nameOrFindFunction);
+            if (expectedFirstDeclarationName == null)
+                expect(statement).to.be.undefined;
+            else
+                expect(statement!.getDeclarations()[0].getName()).to.equal(expectedFirstDeclarationName);
+        }
+
+        it("should get a variable statement when something matches the provided identifier", () => {
+            doTest("Identifier2", "Identifier2");
+        });
+
+        it("should get a variable statement by identifier name when there are multiple", () => {
+            doTest("Identifier3", "Identifier2");
+        });
+
+        it("should return undefined when nothing matches the provided name", () => {
+            doTest("DoesNotExist", undefined);
+        });
+
+        it("should get a variable statement when something matches the provided function", () => {
+            doTest(s => s.getDeclarations().length === 2, "Identifier2");
         });
 
-        it("should return undefined when nothing matches", () => {
-            const statement = variablesSourceFile.getVariableStatement(s => s.getDeclarations().length === 5);
-            expect(statement).to.be.undefined;
+        it("should return undefined when nothing matches the provided function", () => {
+            doTest(s => s.getDeclarations().length === 5, undefined);
         });
     });
 
     describe(nameof<StatementedNode>(n => n.getVariableStatementOrThrow), () => {
-        it("should get a variable statement when something matches", () => {
-            const statement = variablesSourceFile.getVariableStatementOrThrow(s => s.getDeclarations().length === 2);
-            expect(statement.getDeclarations()[0].getName()).to.equal("Identifier2");
+        function doTrueTest(nameOrFindFunction: string | ((declaration: VariableStatement) => boolean), expectedFirstDeclarationName: string) {
+            const statement = variablesSourceFile.getVariableStatementOrThrow(nameOrFindFunction);
+            expect(statement.getDeclarations()[0].getName()).to.equal(expectedFirstDeclarationName);
+        }
+
+        it("should get a variable statement by identifier name", () => {
+            doTrueTest("Identifier1", "Identifier1");
+        });
+
+        it("should get a variable statement by identifier name when there are multiple", () => {
+            doTrueTest("Identifier3", "Identifier2");
+        });
+
+        it("should throw when nothing matches the provided name", () => {
+            expect(() => variablesSourceFile.getVariableStatementOrThrow("DoesNotExist")).to.throw();
+        });
+
+        it("should get a variable statement by function when something matches", () => {
+            doTrueTest(s => s.getDeclarations().length === 2, "Identifier2");
         });
 
-        it("should throw when nothing matches", () => {
+        it("should throw when nothing matches the provided function", () => {
             expect(() => variablesSourceFile.getVariableStatementOrThrow(s => s.getDeclarations().length === 5)).to.throw();
         });
     });