feat: #655 - ExpressionNode - Add getExpressionIfKind methods.

dsherret · Oct 5, 2019 · 332fb7e · 332fb7e
1 parent fda5970
commit 332fb7e
Show file tree

Hide file tree

Showing 3 changed files with 70 additions and 1 deletion.
diff --git a/lib/ts-morph.d.ts b/lib/ts-morph.d.ts
@@ -5623,6 +5623,16 @@ export interface ExpressionedNode {
      * Gets the expression.
      */
     getExpression(): Expression;
+    /**
+     * Gets the expression if its of a certain kind or returns undefined.
+     * @param kind - Syntax kind of the expression.
+     */
+    getExpressionIfKind<TKind extends SyntaxKind>(kind: TKind): KindToExpressionMappings[TKind] | undefined;
+    /**
+     * Gets the expression if its of a certain kind or throws.
+     * @param kind - Syntax kind of the expression.
+     */
+    getExpressionIfKindOrThrow<TKind extends SyntaxKind>(kind: TKind): KindToExpressionMappings[TKind];
     /**
      * Sets the expression.
      * @param textOrWriterFunction - Text to set the expression with.

diff --git a/src/compiler/ast/expression/expressioned/ExpressionedNode.ts b/src/compiler/ast/expression/expressioned/ExpressionedNode.ts
@@ -1,7 +1,10 @@
 import { ExpressionedNodeStructure } from "../../../../structures";
 import { Constructor, WriterFunction } from "../../../../types";
-import { ts } from "../../../../typescript";
+import { ts, SyntaxKind } from "../../../../typescript";
+import * as errors from "../../../../errors";
+import { getSyntaxKindName } from "../../../../utils";
 import { callBaseSet } from "../../callBaseSet";
+import { KindToExpressionMappings } from "../../kindToNodeMappings";
 import { Node } from "../../common";
 import { Expression } from "../Expression";
 
@@ -12,6 +15,16 @@ export interface ExpressionedNode {
      * Gets the expression.
      */
     getExpression(): Expression;
+    /**
+     * Gets the expression if its of a certain kind or returns undefined.
+     * @param kind - Syntax kind of the expression.
+     */
+    getExpressionIfKind<TKind extends SyntaxKind>(kind: TKind): KindToExpressionMappings[TKind] | undefined;
+    /**
+     * Gets the expression if its of a certain kind or throws.
+     * @param kind - Syntax kind of the expression.
+     */
+    getExpressionIfKindOrThrow<TKind extends SyntaxKind>(kind: TKind): KindToExpressionMappings[TKind];
     /**
      * Sets the expression.
      * @param textOrWriterFunction - Text to set the expression with.
@@ -25,6 +38,15 @@ export function ExpressionedNode<T extends Constructor<ExpressionedNodeExtension
             return this._getNodeFromCompilerNode(this.compilerNode.expression);
         }
 
+        getExpressionIfKind<TKind extends SyntaxKind>(kind: TKind): KindToExpressionMappings[TKind] | undefined {
+            const { expression } = this.compilerNode;
+            return expression.kind === kind ? (this._getNodeFromCompilerNode(expression) as KindToExpressionMappings[TKind]) : undefined;
+        }
+
+        getExpressionIfKindOrThrow<TKind extends SyntaxKind>(kind: TKind): KindToExpressionMappings[TKind] {
+            return errors.throwIfNullOrUndefined(this.getExpressionIfKind(kind), `An expression of the kind ${getSyntaxKindName(kind)} was expected.`);
+        }
+
         setExpression(textOrWriterFunction: string | WriterFunction) {
             this.getExpression().replaceWithText(textOrWriterFunction);
             return this;

diff --git a/src/tests/compiler/ast/expression/expressioned/expressionedNodeTests.ts b/src/tests/compiler/ast/expression/expressioned/expressionedNodeTests.ts
@@ -15,6 +15,43 @@ describe(nameof(ExpressionedNode), () => {
         });
     });
 
+    describe(nameof<ExpressionedNode>(n => n.getExpressionIfKind), () => {
+        function doTest(text: string, kind: SyntaxKind, expectedText: string | undefined) {
+            const { descendant } = getInfoFromTextWithDescendant<ParenthesizedExpression>(text, SyntaxKind.ParenthesizedExpression);
+            const result = descendant.getExpressionIfKind(kind);
+            if (expectedText == null)
+                expect(result).to.be.undefined;
+            else
+                expect(result!.getText()).to.equal(expectedText);
+        }
+
+        it("should get the expression when providing the expected value", () => {
+            doTest("(x + 1)", SyntaxKind.BinaryExpression, "x + 1");
+        });
+
+        it("should not get the expression when providing something else", () => {
+            doTest("(x + 1)", SyntaxKind.CallExpression, undefined);
+        });
+    });
+
+    describe(nameof<ExpressionedNode>(n => n.getExpressionIfKindOrThrow), () => {
+        function doTest(text: string, kind: SyntaxKind, expectedText: string | undefined) {
+            const { descendant } = getInfoFromTextWithDescendant<ParenthesizedExpression>(text, SyntaxKind.ParenthesizedExpression);
+            if (expectedText == null)
+                expect(() => descendant.getExpressionIfKindOrThrow(kind)).to.throw();
+            else
+                expect(descendant.getExpressionIfKindOrThrow(kind).getText()).to.equal(expectedText);
+        }
+
+        it("should get the expression when providing the expected value", () => {
+            doTest("(x + 1)", SyntaxKind.BinaryExpression, "x + 1");
+        });
+
+        it("should not get the expression when providing something else", () => {
+            doTest("(x + 1)", SyntaxKind.CallExpression, undefined);
+        });
+    });
+
     describe(nameof<ExpressionedNode>(n => n.setExpression), () => {
         function doTest(text: string, newText: string, expected: string) {
             const { descendant, sourceFile } = getInfoFromTextWithDescendant<ParenthesizedExpression>(text, SyntaxKind.ParenthesizedExpression);