feat: Ability to set a return type with a writer.

dsherret · Apr 30, 2018 · 3db81f8 · 3db81f8
1 parent 5dc3565
commit 3db81f8
Show file tree

Hide file tree

Showing 14 changed files with 73 additions and 18 deletions.
diff --git a/src/compiler/base/ReturnTypedNode.ts b/src/compiler/base/ReturnTypedNode.ts
@@ -1,10 +1,11 @@
 import { ts, SyntaxKind } from "../../typescript";
+import { CodeBlockWriter } from "../../codeBlockWriter";
 import { Constructor } from "../../Constructor";
 import { ReturnTypedNodeStructure } from "../../structures";
 import { callBaseFill } from "../callBaseFill";
 import * as errors from "../../errors";
 import { removeChildren, insertIntoParentTextRange } from "../../manipulation";
-import { StringUtils } from "../../utils";
+import { StringUtils, getTextFromStringOrWriter } from "../../utils";
 import { Node } from "../common";
 import { Type } from "../type/Type";
 import { TypeNode } from "../type/TypeNode";
@@ -24,6 +25,11 @@ export interface ReturnTypedNode {
      * Gets the return type node or throws if none exists.
      */
     getReturnTypeNodeOrThrow(): TypeNode;
+    /**
+     * Sets the return type of the node.
+     * @param writerFunction - Writer function to set the return type with.
+     */
+    setReturnType(writerFunction: (writer: CodeBlockWriter) => void): this;
     /**
      * Sets the return type of the node.
      * @param text - Text to set as the type.
@@ -51,7 +57,8 @@ export function ReturnTypedNode<T extends Constructor<ReturnTypedNodeExtensionRe
             return errors.throwIfNullOrUndefined(this.getReturnTypeNode(), "Expected to find a return type node.");
         }
 
-        setReturnType(text: string) {
+        setReturnType(textOrWriterFunction: string | ((writer: CodeBlockWriter) => void)) {
+            const text = getTextFromStringOrWriter(this.getWriterWithQueuedChildIndentation(), textOrWriterFunction);
             if (StringUtils.isNullOrWhitespace(text))
                 return this.removeReturnType();
 

diff --git a/src/factories/StructurePrinterFactory.ts b/src/factories/StructurePrinterFactory.ts
@@ -24,6 +24,11 @@ export class StructurePrinterFactory {
         return new structurePrinters.ModifierableNodeStructurePrinter(this);
     }
 
+    @Memoize
+    forReturnTypedNode(alwaysWrite?: boolean): structurePrinters.ReturnTypedNodeStructurePrinter {
+        return new structurePrinters.ReturnTypedNodeStructurePrinter(this, alwaysWrite);
+    }
+
     @Memoize
     forTypedNode(separator: string, alwaysWrite?: boolean): structurePrinters.TypedNodeStructurePrinter {
         return new structurePrinters.TypedNodeStructurePrinter(this, separator, alwaysWrite);

diff --git a/src/structurePrinters/base/ReturnTypedNodeStructurePrinter.ts b/src/structurePrinters/base/ReturnTypedNodeStructurePrinter.ts
@@ -0,0 +1,30 @@
+import { CodeBlockWriter } from "../../codeBlockWriter";
+import { ReturnTypedNodeStructure } from "../../structures";
+import { StructurePrinterFactory } from "../../factories";
+import { StringUtils } from "../../utils";
+import { FactoryStructurePrinter } from "../FactoryStructurePrinter";
+
+export class ReturnTypedNodeStructurePrinter extends FactoryStructurePrinter<ReturnTypedNodeStructure> {
+    constructor(factory: StructurePrinterFactory, private readonly alwaysWrite = false) {
+        super(factory);
+    }
+
+    printText(writer: CodeBlockWriter, structure: ReturnTypedNodeStructure) {
+        let { returnType } = structure;
+        if (returnType == null && this.alwaysWrite === false)
+            return;
+
+        returnType = returnType || "void";
+
+        // todo: hacky, will need to change this in the future...
+        const initializerText = typeof returnType === "string" ? returnType : getTextForWriterFunc(returnType);
+        if (!StringUtils.isNullOrWhitespace(initializerText))
+            writer.write(`: ${initializerText}`);
+
+        function getTextForWriterFunc(writerFunc: (writer: CodeBlockWriter) => void) {
+            const newWriter = new CodeBlockWriter(writer.getOptions());
+            writerFunc(newWriter);
+            return newWriter.toString();
+        }
+    }
+}
diff --git a/src/structurePrinters/base/TypedNodeStructurePrinter.ts b/src/structurePrinters/base/TypedNodeStructurePrinter.ts
@@ -1,5 +1,5 @@
 import { CodeBlockWriter } from "../../codeBlockWriter";
-import { TypedNodeStructure, QuestionTokenableNodeStructure } from "../../structures";
+import { TypedNodeStructure } from "../../structures";
 import { StructurePrinterFactory } from "../../factories";
 import { StringUtils } from "../../utils";
 import { FactoryStructurePrinter } from "../FactoryStructurePrinter";

diff --git a/src/structurePrinters/base/index.ts b/src/structurePrinters/base/index.ts
@@ -1,3 +1,4 @@
 export * from "./InitializerExpressionableNodeStructurePrinter";
 export * from "./ModifierableNodeStructurePrinter";
+export * from "./ReturnTypedNodeStructurePrinter";
 export * from "./TypedNodeStructurePrinter";
diff --git a/src/structurePrinters/class/GetAccessorDeclarationStructurePrinter.ts b/src/structurePrinters/class/GetAccessorDeclarationStructurePrinter.ts
@@ -1,6 +1,5 @@
 import { CodeBlockWriter } from "../../codeBlockWriter";
 import { GetAccessorDeclarationStructure } from "../../structures";
-import { StringUtils } from "../../utils";
 import { StructurePrinterFactory } from "../../factories";
 import { FactoryStructurePrinter } from "../FactoryStructurePrinter";
 import { BlankLineFormattingStructuresPrinter } from "../formatting";
@@ -25,7 +24,7 @@ export class GetAccessorDeclarationStructurePrinter extends FactoryStructurePrin
         writer.write("(");
         this.factory.forParameterDeclaration().printTexts(writer, structure.parameters);
         writer.write(")");
-        writer.conditionalWrite(!StringUtils.isNullOrWhitespace(structure.returnType), `: ${structure.returnType}`);
+        this.factory.forReturnTypedNode().printText(writer, structure);
         writer.spaceIfLastNot().inlineBlock(() => {
             this.factory.forBodyText(this.options).printText(writer, structure);
         });

diff --git a/src/structurePrinters/class/MethodDeclarationStructurePrinter.ts b/src/structurePrinters/class/MethodDeclarationStructurePrinter.ts
@@ -75,6 +75,6 @@ export class MethodDeclarationStructurePrinter extends FactoryStructurePrinter<M
         writer.write("(");
         this.factory.forParameterDeclaration().printTexts(writer, structure.parameters);
         writer.write(`)`);
-        writer.conditionalWrite(structure.returnType != null && structure.returnType.length > 0, `: ${structure.returnType}`);
+        this.factory.forReturnTypedNode().printText(writer, structure);
     }
 }
diff --git a/src/structurePrinters/class/SetAccessorDeclarationStructurePrinter.ts b/src/structurePrinters/class/SetAccessorDeclarationStructurePrinter.ts
@@ -1,6 +1,5 @@
 import { CodeBlockWriter } from "../../codeBlockWriter";
 import { SetAccessorDeclarationStructure } from "../../structures";
-import { StringUtils } from "../../utils";
 import { StructurePrinterFactory } from "../../factories";
 import { FactoryStructurePrinter } from "../FactoryStructurePrinter";
 import { BlankLineFormattingStructuresPrinter } from "../formatting";
@@ -25,7 +24,7 @@ export class SetAccessorDeclarationStructurePrinter extends FactoryStructurePrin
         writer.write("(");
         this.factory.forParameterDeclaration().printTexts(writer, structure.parameters);
         writer.write(")");
-        writer.conditionalWrite(!StringUtils.isNullOrWhitespace(structure.returnType), `: ${structure.returnType}`);
+        this.factory.forReturnTypedNode().printText(writer, structure);
         writer.spaceIfLastNot().inlineBlock(() => {
             this.factory.forBodyText(this.options).printText(writer, structure);
         });

diff --git a/src/structurePrinters/function/FunctionDeclarationStructurePrinter.ts b/src/structurePrinters/function/FunctionDeclarationStructurePrinter.ts
@@ -1,6 +1,6 @@
 import { CodeBlockWriter } from "../../codeBlockWriter";
 import { FunctionDeclarationStructure, FunctionDeclarationOverloadStructure } from "../../structures";
-import { StringUtils, setValueIfUndefined, ObjectUtils } from "../../utils";
+import { setValueIfUndefined, ObjectUtils } from "../../utils";
 import { StructurePrinterFactory } from "../../factories";
 import { FactoryStructurePrinter } from "../FactoryStructurePrinter";
 
@@ -75,7 +75,6 @@ export class FunctionDeclarationStructurePrinter extends FactoryStructurePrinter
         if (structure.parameters != null)
             this.factory.forParameterDeclaration().printTexts(writer, structure.parameters);
         writer.write(`)`);
-        if (!StringUtils.isNullOrWhitespace(structure.returnType))
-            writer.write(`: ${structure.returnType}`);
+        this.factory.forReturnTypedNode().printText(writer, structure);
     }
 }
diff --git a/src/structurePrinters/interface/CallSignatureDeclarationStructurePrinter.ts b/src/structurePrinters/interface/CallSignatureDeclarationStructurePrinter.ts
@@ -16,7 +16,7 @@ export class CallSignatureDeclarationStructurePrinter extends FactoryStructurePr
         writer.write("(");
         this.factory.forParameterDeclaration().printTexts(writer, structure.parameters);
         writer.write(")");
-        writer.write(`: ${structure.returnType || "void"}`);
+        this.factory.forReturnTypedNode(true).printText(writer, structure);
         writer.write(";");
     }
 }
diff --git a/src/structurePrinters/interface/ConstructSignatureDeclarationStructurePrinter.ts b/src/structurePrinters/interface/ConstructSignatureDeclarationStructurePrinter.ts
@@ -17,8 +17,7 @@ export class ConstructSignatureDeclarationStructurePrinter extends FactoryStruct
         writer.write("(");
         this.factory.forParameterDeclaration().printTexts(writer, structure.parameters);
         writer.write(")");
-        if (structure.returnType != null && structure.returnType.length > 0)
-            writer.write(`: ${structure.returnType}`);
+        this.factory.forReturnTypedNode().printText(writer, structure);
         writer.write(";");
     }
 }
diff --git a/src/structurePrinters/interface/MethodSignatureStructurePrinter.ts b/src/structurePrinters/interface/MethodSignatureStructurePrinter.ts
@@ -18,8 +18,7 @@ export class MethodSignatureStructurePrinter extends FactoryStructurePrinter<Met
         writer.write("(");
         this.factory.forParameterDeclaration().printTexts(writer, structure.parameters);
         writer.write(")");
-        if (structure.returnType != null && structure.returnType.length > 0)
-            writer.write(`: ${structure.returnType}`);
+        this.factory.forReturnTypedNode().printText(writer, structure);
         writer.write(";");
     }
 }
diff --git a/src/structures/base/ReturnTypedNodeStructure.ts b/src/structures/base/ReturnTypedNodeStructure.ts
@@ -1,3 +1,5 @@
-export interface ReturnTypedNodeStructure {
-    returnType?: string;
+import { CodeBlockWriter } from "../../codeBlockWriter";
+
+export interface ReturnTypedNodeStructure {
+    returnType?: string | ((writer: CodeBlockWriter) => void);
 }
diff --git a/src/tests/compiler/base/returnTypedNodeTests.ts b/src/tests/compiler/base/returnTypedNodeTests.ts
@@ -1,4 +1,5 @@
 import { expect } from "chai";
+import { CodeBlockWriter } from "../../../codeBlockWriter";
 import { ReturnTypedNode, FunctionDeclaration } from "../../../compiler";
 import { ReturnTypedNodeStructure } from "../../../structures";
 import { getInfoFromText } from "../testHelpers";
@@ -64,6 +65,16 @@ describe(nameof(ReturnTypedNode), () => {
         it("should remove the type when it's empty", () => {
             doTest("function identifier(): string {}", "", "function identifier() {}");
         });
+
+        function doWriterTest(startText: string, returnType: (writer: CodeBlockWriter) => void, expectedText: string) {
+            const {firstChild} = getInfoFromText<FunctionDeclaration>(startText);
+            firstChild.setReturnType(returnType);
+            expect(firstChild.getText()).to.equal(expectedText);
+        }
+
+        it("should set the return type with a writer", () => {
+            doWriterTest("declare function identifier(): number;", writer => writer.write("string"), "declare function identifier(): string;");
+        });
     });
 
     describe(nameof<ReturnTypedNode>(n => n.removeReturnType), () => {
@@ -93,6 +104,10 @@ describe(nameof(ReturnTypedNode), () => {
             doTest("function Identifier() {}", { returnType: "number" }, "function Identifier(): number {}");
         });
 
+        it("should modify when setting with a writer", () => {
+            doTest("function Identifier() {}", { returnType: writer => writer.write("number") }, "function Identifier(): number {}");
+        });
+
         it("should not modify anything if the structure doesn't change anything", () => {
             doTest("function Identifier() {}", { }, "function Identifier() {}");
         });