#27 - Adding/inserting arguments on decorators and call expressions.

dsherret · Nov 12, 2017 · bf75fda · bf75fda
1 parent 05a18c1
commit bf75fda
Show file tree

Hide file tree

Showing 5 changed files with 249 additions and 8 deletions.
diff --git a/src/compiler/base/ArgumentedNode.ts b/src/compiler/base/ArgumentedNode.ts
@@ -1,7 +1,8 @@
 import * as ts from "typescript";
 import * as errors from "./../../errors";
 import {Constructor} from "./../../Constructor";
-import {removeCommaSeparatedChild, verifyAndGetIndex} from "./../../manipulation";
+import {removeCommaSeparatedChild, verifyAndGetIndex, insertIntoCommaSeparatedNodes} from "./../../manipulation";
+import {ArrayUtils} from "./../../utils";
 import {Node} from "./../common";
 
 export type ArgumentedNodeExtensionType = Node<ts.Node & { arguments: ts.NodeArray<ts.Node>; }>;
@@ -11,6 +12,28 @@ export interface ArgumentedNode {
      * Gets all the arguments of the node.
      */
     getArguments(): Node[];
+    /**
+     * Adds an argument.
+     * @param argumentText - Argument text to add.
+     */
+    addArgument(argumentText: string): Node;
+    /**
+     * Adds arguments.
+     * @param argumentTexts - Argument texts to add.
+     */
+    addArguments(argumentTexts: string[]): Node[];
+    /**
+     * Inserts an argument.
+     * @param index - Index to insert at.
+     * @param argumentText - Argument text to insert.
+     */
+    insertArgument(index: number, argumentText: string): Node;
+    /**
+     * Inserts arguments.
+     * @param index - Index to insert at.
+     * @param argumentTexts - Argument texts to insert.
+     */
+    insertArguments(index: number, argumentTexts: string[]): Node[];
     /**
      * Removes an argument.
      * @param arg - Argument to remove.
@@ -22,9 +45,7 @@ export interface ArgumentedNode {
      */
     removeArgument(index: number): this;
     /**
-     * Removes an argument argument.
      * @internal
-     * @param argOrIndex - Argument of index to remove.
      */
     removeArgument(argOrIndex: Node | number): this;
 }
@@ -35,6 +56,35 @@ export function ArgumentedNode<T extends Constructor<ArgumentedNodeExtensionType
             return this.compilerNode.arguments.map(a => this.global.compilerFactory.getNodeFromCompilerNode(a, this.sourceFile)) as Node[];
         }
 
+        addArgument(argumentText: string) {
+            return this.addArguments([argumentText])[0];
+        }
+
+        addArguments(argumentTexts: string[]) {
+            return this.insertArguments(this.getArguments().length, argumentTexts);
+        }
+
+        insertArgument(index: number, argumentText: string) {
+            return this.insertArguments(index, [argumentText])[0];
+        }
+
+        insertArguments(index: number, argumentTexts: string[]) {
+            if (ArrayUtils.isNullOrEmpty(argumentTexts))
+                return [];
+
+            const args = this.getArguments();
+            index = verifyAndGetIndex(index, args.length);
+
+            insertIntoCommaSeparatedNodes({
+                parent: this,
+                currentNodes: args,
+                insertIndex: index,
+                newTexts: argumentTexts
+            });
+
+            return this.getArguments().slice(index, index + argumentTexts.length);
+        }
+
         removeArgument(arg: Node): this;
         removeArgument(index: number): this;
         removeArgument(argOrIndex: Node | number): this {

diff --git a/src/compiler/decorator/Decorator.ts b/src/compiler/decorator/Decorator.ts
@@ -209,6 +209,41 @@ export class Decorator extends DecoratorBase<ts.Decorator> {
         return this;
     }
 
+    /**
+     * Adds an argument.
+     * @param argumentTexts - Argument text.
+     */
+    addArgument(argumentText: string) {
+        return this.addArguments([argumentText])[0];
+    }
+
+    /**
+     * Adds arguments.
+     * @param argumentTexts - Argument texts.
+     */
+    addArguments(argumentTexts: string[]) {
+        return this.insertArguments(this.getArguments().length, argumentTexts);
+    }
+
+    /**
+     * Inserts an argument.
+     * @param index - Index to insert at.
+     * @param argumentTexts - Argument text.
+     */
+    insertArgument(index: number, argumentText: string) {
+        return this.insertArguments(index, [argumentText])[0];
+    }
+
+    /**
+     * Inserts arguments.
+     * @param index - Index to insert at.
+     * @param argumentTexts - Argument texts.
+     */
+    insertArguments(index: number, argumentTexts: string[]) {
+        this.setIsDecoratorFactory(true);
+        return this.getCallExpressionOrThrow().insertArguments(index, argumentTexts);
+    }
+
     /**
      * Removes an argument based on the node.
      * @param node - Argument's node to remove.

diff --git a/src/tests/compiler/base/argumentedNodeTests.ts b/src/tests/compiler/base/argumentedNodeTests.ts
@@ -19,6 +19,78 @@ describe(nameof(ArgumentedNode), () => {
         });
     });
 
+    describe(nameof<ArgumentedNode>(n => n.insertArguments), () => {
+        function doTest(code: string, index: number, texts: string[], expectedCode: string) {
+            const {firstChild, sourceFile} = getInfoFromText<ClassDeclaration>(code);
+            const callExpr = firstChild.getDecorators()[0].getCallExpressionOrThrow();
+            const result = callExpr.insertArguments(index, texts);
+            expect(result.map(t => t.getText())).to.deep.equal(texts);
+            expect(sourceFile.getFullText()).to.equal(expectedCode);
+        }
+
+        it("should insert multiple args when none exist", () => {
+            doTest("@dec()\nclass T {}", 0, ["5", "6", "7"], "@dec(5, 6, 7)\nclass T {}");
+        });
+
+        it("should insert multiple args at the beginning", () => {
+            doTest("@dec(3)\nclass T {}", 0, ["1", "2"], "@dec(1, 2, 3)\nclass T {}");
+        });
+
+        it("should insert multiple args in the middle", () => {
+            doTest("@dec(1, 4)\nclass T {}", 1, ["2", "3"], "@dec(1, 2, 3, 4)\nclass T {}");
+        });
+
+        it("should insert multiple args at the end", () => {
+            doTest("@dec(1)\nclass T {}", 1, ["2", "3"], "@dec(1, 2, 3)\nclass T {}");
+        });
+
+        it("should insert args when a type argument exists", () => {
+            doTest("@dec<3>(1)\nclass T {}", 1, ["2", "3"], "@dec<3>(1, 2, 3)\nclass T {}");
+        });
+    });
+
+    describe(nameof<ArgumentedNode>(n => n.insertArgument), () => {
+        function doTest(code: string, index: number, text: string, expectedCode: string) {
+            const {firstChild, sourceFile} = getInfoFromText<ClassDeclaration>(code);
+            const callExpr = firstChild.getDecorators()[0].getCallExpressionOrThrow();
+            const result = callExpr.insertArgument(index, text);
+            expect(result.getText()).to.equal(text);
+            expect(sourceFile.getFullText()).to.equal(expectedCode);
+        }
+
+        it("should insert an arg", () => {
+            doTest("@dec(1, 3)\nclass T {}", 1, "2", "@dec(1, 2, 3)\nclass T {}");
+        });
+    });
+
+    describe(nameof<ArgumentedNode>(n => n.addArguments), () => {
+        function doTest(code: string, texts: string[], expectedCode: string) {
+            const {firstChild, sourceFile} = getInfoFromText<ClassDeclaration>(code);
+            const callExpr = firstChild.getDecorators()[0].getCallExpressionOrThrow();
+            const result = callExpr.addArguments(texts);
+            expect(result.map(t => t.getText())).to.deep.equal(texts);
+            expect(sourceFile.getFullText()).to.equal(expectedCode);
+        }
+
+        it("should add multiple args", () => {
+            doTest("@dec(1)\nclass T {}", ["2", "3"], "@dec(1, 2, 3)\nclass T {}");
+        });
+    });
+
+    describe(nameof<ArgumentedNode>(n => n.addArgument), () => {
+        function doTest(code: string, text: string, expectedCode: string) {
+            const {firstChild, sourceFile} = getInfoFromText<ClassDeclaration>(code);
+            const callExpr = firstChild.getDecorators()[0].getCallExpressionOrThrow();
+            const result = callExpr.addArgument(text);
+            expect(result.getText()).to.equal(text);
+            expect(sourceFile.getFullText()).to.equal(expectedCode);
+        }
+
+        it("should add an arg", () => {
+            doTest("@dec(1, 2)\nclass T {}", "3", "@dec(1, 2, 3)\nclass T {}");
+        });
+    });
+
     describe(nameof<ArgumentedNode>(d => d.removeArgument), () => {
         function doTest(text: string, removeIndex: number, expectedText: string) {
             doTestByIndex();

diff --git a/src/tests/compiler/base/typeArgumentedNodeTests.ts b/src/tests/compiler/base/typeArgumentedNodeTests.ts
@@ -6,7 +6,7 @@ describe(nameof(TypeArgumentedNode), () => {
     describe(nameof<TypeArgumentedNode>(n => n.getTypeArguments), () => {
         function doTest(code: string, expectedArgs: string[]) {
             const {firstChild} = getInfoFromText<ClassDeclaration>(code);
-            const args = firstChild.getDecorators()[0].getCallExpression()!.getTypeArguments();
+            const args = firstChild.getDecorators()[0].getCallExpressionOrThrow().getTypeArguments();
             expect(args.map(a => a.getText())).to.deep.equal(expectedArgs);
         }
 
@@ -22,7 +22,7 @@ describe(nameof(TypeArgumentedNode), () => {
     describe(nameof<TypeArgumentedNode>(n => n.insertTypeArguments), () => {
         function doTest(code: string, index: number, texts: string[], expectedCode: string) {
             const {firstChild, sourceFile} = getInfoFromText<ClassDeclaration>(code);
-            const callExpr = firstChild.getDecorators()[0].getCallExpression()!;
+            const callExpr = firstChild.getDecorators()[0].getCallExpressionOrThrow();
             const result = callExpr.insertTypeArguments(index, texts);
             expect(result.map(t => t.getText())).to.deep.equal(texts);
             expect(sourceFile.getFullText()).to.equal(expectedCode);
@@ -48,7 +48,7 @@ describe(nameof(TypeArgumentedNode), () => {
     describe(nameof<TypeArgumentedNode>(n => n.insertTypeArgument), () => {
         function doTest(code: string, index: number, text: string, expectedCode: string) {
             const {firstChild, sourceFile} = getInfoFromText<ClassDeclaration>(code);
-            const callExpr = firstChild.getDecorators()[0].getCallExpression()!;
+            const callExpr = firstChild.getDecorators()[0].getCallExpressionOrThrow();
             const result = callExpr.insertTypeArgument(index, text);
             expect(result.getText()).to.equal(text);
             expect(sourceFile.getFullText()).to.equal(expectedCode);
@@ -62,7 +62,7 @@ describe(nameof(TypeArgumentedNode), () => {
     describe(nameof<TypeArgumentedNode>(n => n.addTypeArguments), () => {
         function doTest(code: string, texts: string[], expectedCode: string) {
             const {firstChild, sourceFile} = getInfoFromText<ClassDeclaration>(code);
-            const callExpr = firstChild.getDecorators()[0].getCallExpression()!;
+            const callExpr = firstChild.getDecorators()[0].getCallExpressionOrThrow();
             const result = callExpr.addTypeArguments(texts);
             expect(result.map(t => t.getText())).to.deep.equal(texts);
             expect(sourceFile.getFullText()).to.equal(expectedCode);
@@ -76,7 +76,7 @@ describe(nameof(TypeArgumentedNode), () => {
     describe(nameof<TypeArgumentedNode>(n => n.addTypeArgument), () => {
         function doTest(code: string, text: string, expectedCode: string) {
             const {firstChild, sourceFile} = getInfoFromText<ClassDeclaration>(code);
-            const callExpr = firstChild.getDecorators()[0].getCallExpression()!;
+            const callExpr = firstChild.getDecorators()[0].getCallExpressionOrThrow();
             const result = callExpr.addTypeArgument(text);
             expect(result.getText()).to.equal(text);
             expect(sourceFile.getFullText()).to.equal(expectedCode);

diff --git a/src/tests/compiler/decorator/decoratorTests.ts b/src/tests/compiler/decorator/decoratorTests.ts
@@ -165,6 +165,90 @@ describe(nameof(Decorator), () => {
         });
     });
 
+    describe(nameof<Decorator>(n => n.insertArguments), () => {
+        function doTest(code: string, index: number, texts: string[], expectedCode: string) {
+            const {firstChild, sourceFile} = getInfoFromText<ClassDeclaration>(code);
+            const result = firstChild.getDecorators()[0].insertArguments(index, texts);
+            expect(result.map(t => t.getText())).to.deep.equal(texts);
+            expect(sourceFile.getFullText()).to.equal(expectedCode);
+        }
+
+        it("should insert multiple args when none exist", () => {
+            doTest("@dec()\nclass T {}", 0, ["5", "6", "7"], "@dec(5, 6, 7)\nclass T {}");
+        });
+
+        it("should insert multiple args at the beginning", () => {
+            doTest("@dec(3)\nclass T {}", 0, ["1", "2"], "@dec(1, 2, 3)\nclass T {}");
+        });
+
+        it("should insert multiple args in the middle", () => {
+            doTest("@dec(1, 4)\nclass T {}", 1, ["2", "3"], "@dec(1, 2, 3, 4)\nclass T {}");
+        });
+
+        it("should insert multiple args at the end", () => {
+            doTest("@dec(1)\nclass T {}", 1, ["2", "3"], "@dec(1, 2, 3)\nclass T {}");
+        });
+
+        it("should insert args when a type argument exists", () => {
+            doTest("@dec<3>(1)\nclass T {}", 1, ["2", "3"], "@dec<3>(1, 2, 3)\nclass T {}");
+        });
+
+        it("should set as decorator factory when not", () => {
+            doTest("@dec\nclass T {}", 0, ["1"], "@dec(1)\nclass T {}");
+        });
+    });
+
+    describe(nameof<Decorator>(n => n.insertArgument), () => {
+        function doTest(code: string, index: number, text: string, expectedCode: string) {
+            const {firstChild, sourceFile} = getInfoFromText<ClassDeclaration>(code);
+            const result = firstChild.getDecorators()[0].insertArgument(index, text);
+            expect(result.getText()).to.equal(text);
+            expect(sourceFile.getFullText()).to.equal(expectedCode);
+        }
+
+        it("should insert an arg", () => {
+            doTest("@dec(1, 3)\nclass T {}", 1, "2", "@dec(1, 2, 3)\nclass T {}");
+        });
+
+        it("should set as decorator factory when not", () => {
+            doTest("@dec\nclass T {}", 0, "1", "@dec(1)\nclass T {}");
+        });
+    });
+
+    describe(nameof<Decorator>(n => n.addArguments), () => {
+        function doTest(code: string, texts: string[], expectedCode: string) {
+            const {firstChild, sourceFile} = getInfoFromText<ClassDeclaration>(code);
+            const result = firstChild.getDecorators()[0].addArguments(texts);
+            expect(result.map(t => t.getText())).to.deep.equal(texts);
+            expect(sourceFile.getFullText()).to.equal(expectedCode);
+        }
+
+        it("should add multiple args", () => {
+            doTest("@dec(1)\nclass T {}", ["2", "3"], "@dec(1, 2, 3)\nclass T {}");
+        });
+
+        it("should set as decorator factory when not", () => {
+            doTest("@dec\nclass T {}", ["1"], "@dec(1)\nclass T {}");
+        });
+    });
+
+    describe(nameof<Decorator>(n => n.addArgument), () => {
+        function doTest(code: string, text: string, expectedCode: string) {
+            const {firstChild, sourceFile} = getInfoFromText<ClassDeclaration>(code);
+            const result = firstChild.getDecorators()[0].addArgument(text);
+            expect(result.getText()).to.equal(text);
+            expect(sourceFile.getFullText()).to.equal(expectedCode);
+        }
+
+        it("should add an arg", () => {
+            doTest("@dec(1, 2)\nclass T {}", "3", "@dec(1, 2, 3)\nclass T {}");
+        });
+
+        it("should set as decorator factory when not", () => {
+            doTest("@dec\nclass T {}", "1", "@dec(1)\nclass T {}");
+        });
+    });
+
     describe(nameof<Decorator>(d => d.getTypeArguments), () => {
         it("should return an empty array when not a decorator factory", () => {
             const {firstDecorator} = getFirstClassDecorator("@decorator\nclass Identifier {}");