feat: Ability to set an initializer's text using a writer.

dsherret · Apr 29, 2018 · 2c1a9e5 · 2c1a9e5
1 parent 67a22b9
commit 2c1a9e5
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 3 deletions.
diff --git a/src/compiler/base/initializer/InitializerSetExpressionableNode.ts b/src/compiler/base/initializer/InitializerSetExpressionableNode.ts
@@ -1,9 +1,11 @@
+import CodeBlockWriter from "code-block-writer";
 import { ts, SyntaxKind } from "../../../typescript";
 import { Constructor } from "../../../Constructor";
 import * as errors from "../../../errors";
 import { InitializerSetExpressionableNodeStructure } from "../../../structures";
 import { callBaseFill } from "../../callBaseFill";
 import { insertIntoParentTextRange, removeChildren } from "../../../manipulation";
+import { getTextFromStringOrWriter } from "../../../utils";
 import { Expression } from "../../expression";
 import { Node } from "../../common";
 import { InitializerGetExpressionableNode } from "./InitializerGetExpressionableNode";
@@ -20,6 +22,11 @@ export interface InitializerSetExpressionableNode {
      * @param text - New text to set for the initializer.
      */
     setInitializer(text: string): this;
+    /**
+     * Sets the initializer using a writer function.
+     * @param writerFunction - Function to write the initializer with.
+     */
+    setInitializer(writerFunction: (writer: CodeBlockWriter) => void): this;
 }
 
 export function InitializerSetExpressionableNode<T extends Constructor<InitializerSetExpressionableExtensionType>>(Base: T): Constructor<InitializerSetExpressionableNode> & T {
@@ -37,8 +44,9 @@ export function InitializerSetExpressionableNode<T extends Constructor<Initializ
             return this;
         }
 
-        setInitializer(text: string) {
-            errors.throwIfNotStringOrWhitespace(text, nameof(text));
+        setInitializer(textOrWriterFunction: string | ((writer: CodeBlockWriter) => void)) {
+            const text = getTextFromStringOrWriter(this.getWriter(), textOrWriterFunction);
+            errors.throwIfNotStringOrWhitespace(text, nameof(textOrWriterFunction));
 
             if (this.hasInitializer())
                 this.removeInitializer();

diff --git a/src/tests/compiler/base/initializer/initializerSetExpressionableNodeTests.ts b/src/tests/compiler/base/initializer/initializerSetExpressionableNodeTests.ts
@@ -1,4 +1,5 @@
-import { expect } from "chai";
+import CodeBlockWriter from "code-block-writer";
+import { expect } from "chai";
 import { EnumDeclaration, InitializerSetExpressionableNode, ClassDeclaration, PropertyDeclaration } from "../../../../compiler";
 import { InitializerSetExpressionableNodeStructure } from "../../../../structures";
 import { getInfoFromText } from "../../testHelpers";
@@ -95,6 +96,19 @@ describe(nameof(InitializerSetExpressionableNode), () => {
                 doClassPropTest("class Identifier { prop/*comment*/ }", "2", "class Identifier { prop = 2/*comment*/ }");
             });
         });
+
+        describe("writer", () => {
+            function doClassPropTest(text: string, newInitializer: (writer: CodeBlockWriter) => void, expected: string) {
+                const { firstChild } = getInfoFromText<ClassDeclaration>(text);
+                const prop = firstChild.getInstanceProperties()[0] as PropertyDeclaration;
+                prop.setInitializer(newInitializer);
+                expect(firstChild.getFullText()).to.equal(expected);
+            }
+
+            it("should set a new initializer using a writer", () => {
+                doClassPropTest("class Identifier { prop; }", writer => writer.write("2"), "class Identifier { prop = 2; }");
+            });
+        });
     });
 
     describe(nameof<ClassDeclaration>(n => n.fill), () => {