feat: #605 - Add ability to insert comments on enum declarations.

dsherret · Jun 18, 2019 · 0420ae1 · 0420ae1
1 parent c3b9f28
commit 0420ae1
Show file tree

Hide file tree

Showing 16 changed files with 323 additions and 67 deletions.
diff --git a/lib/code-block-writer.d.ts b/lib/code-block-writer.d.ts
@@ -235,6 +235,11 @@ declare class CodeBlockWriter {
      * Gets the last char written.
      */
     getLastChar(): string | undefined;
+    /**
+     * Gets if the writer ends with the provided text.
+     * @param text - Text to check if the writer ends with the provided text.
+     */
+    endsWith(text: string): boolean;
     /**
      * Iterates over the writer characters in reverse order. The iteration stops when a non-null or
      * undefined value is returned from the action. The returned value is then returned by the method.

diff --git a/lib/ts-morph.d.ts b/lib/ts-morph.d.ts
@@ -5290,23 +5290,45 @@ export declare class EnumDeclaration extends EnumDeclarationBase<ts.EnumDeclarat
      * @param structure - Structure of the enum.
      */
     addMember(structure: OptionalKind<EnumMemberStructure>): EnumMember;
+    /**
+     * Adds a member to the enum.
+     * @param structure - Structure of the enum.
+     */
+    addMember(structure: OptionalKind<EnumMemberStructure> | WriterFunction): EnumMember | CommentEnumMember;
     /**
      * Adds members to the enum.
      * @param structures - Structures of the enums.
      */
     addMembers(structures: ReadonlyArray<OptionalKind<EnumMemberStructure>>): EnumMember[];
+    /**
+     * Adds members to the enum.
+     * @param structures - Structures of the enums.
+     */
+    addMembers(structures: ReadonlyArray<OptionalKind<EnumMemberStructure> | WriterFunction>): (EnumMember | CommentEnumMember)[];
     /**
      * Inserts a member to the enum.
      * @param index - Child index to insert at.
      * @param structure - Structure of the enum.
      */
     insertMember(index: number, structure: OptionalKind<EnumMemberStructure>): EnumMember;
+    /**
+     * Inserts a member to the enum.
+     * @param index - Child index to insert at.
+     * @param structure - Structure of the enum.
+     */
+    insertMember(index: number, structure: OptionalKind<EnumMemberStructure> | WriterFunction): EnumMember | CommentEnumMember;
     /**
      * Inserts members to an enum.
      * @param index - Child index to insert at.
      * @param structures - Structures of the enums.
      */
     insertMembers(index: number, structures: ReadonlyArray<OptionalKind<EnumMemberStructure>>): EnumMember[];
+    /**
+     * Inserts members to an enum.
+     * @param index - Child index to insert at.
+     * @param structures - Structures of the enums.
+     */
+    insertMembers(index: number, structures: ReadonlyArray<OptionalKind<EnumMemberStructure> | WriterFunction>): (EnumMember | CommentEnumMember)[];
     /**
      * Gets an enum member.
      * @param name - Name of the member.

diff --git a/package.json b/package.json
@@ -74,7 +74,7 @@
   },
   "dependencies": {
     "@dsherret/to-absolute-glob": "^2.0.2",
-    "code-block-writer": "9.2.1",
+    "code-block-writer": "9.3.0",
     "fs-extra": "^7.0.0",
     "glob-parent": "^3.1.0",
     "globby": "^8.0.1",

diff --git a/src/codeBlockWriter/code-block-writer.ts b/src/codeBlockWriter/code-block-writer.ts
@@ -238,6 +238,11 @@ declare class CodeBlockWriter {
      * Gets the last char written.
      */
     getLastChar(): string | undefined;
+    /**
+     * Gets if the writer ends with the provided text.
+     * @param text - Text to check if the writer ends with the provided text.
+     */
+    endsWith(text: string): boolean;
     /**
      * Iterates over the writer characters in reverse order. The iteration stops when a non-null or
      * undefined value is returned from the action. The returned value is then returned by the method.

diff --git a/src/compiler/ast/enum/EnumDeclaration.ts b/src/compiler/ast/enum/EnumDeclaration.ts
@@ -11,6 +11,7 @@ import { EnumMember } from "./EnumMember";
 import { callBaseGetStructure } from "../callBaseGetStructure";
 import { CommentEnumMember } from "./CommentEnumMember";
 import { ExtendedParser } from "../utils";
+import { WriterFunction } from "../../../types";
 
 export const EnumDeclarationBase = TextInsertableNode(NamespaceChildableNode(JSDocableNode(AmbientableNode(ExportableNode(
     ModifierableNode(NamedNode(Statement))
@@ -37,15 +38,27 @@ export class EnumDeclaration extends EnumDeclarationBase<ts.EnumDeclaration> {
      * Adds a member to the enum.
      * @param structure - Structure of the enum.
      */
-    addMember(structure: OptionalKind<EnumMemberStructure>) {
+    addMember(structure: OptionalKind<EnumMemberStructure>): EnumMember;
+    /**
+     * Adds a member to the enum.
+     * @param structure - Structure of the enum.
+     */
+    addMember(structure: OptionalKind<EnumMemberStructure> | WriterFunction): EnumMember | CommentEnumMember;
+    addMember(structure: OptionalKind<EnumMemberStructure> | WriterFunction) {
         return this.addMembers([structure])[0];
     }
 
     /**
      * Adds members to the enum.
      * @param structures - Structures of the enums.
      */
-    addMembers(structures: ReadonlyArray<OptionalKind<EnumMemberStructure>>) {
+    addMembers(structures: ReadonlyArray<OptionalKind<EnumMemberStructure>>): EnumMember[];
+    /**
+     * Adds members to the enum.
+     * @param structures - Structures of the enums.
+     */
+    addMembers(structures: ReadonlyArray<OptionalKind<EnumMemberStructure> | WriterFunction>): (EnumMember | CommentEnumMember)[];
+    addMembers(structures: ReadonlyArray<OptionalKind<EnumMemberStructure> | WriterFunction>) {
         return this.insertMembers(this.getMembers().length, structures);
     }
 
@@ -54,7 +67,14 @@ export class EnumDeclaration extends EnumDeclarationBase<ts.EnumDeclaration> {
      * @param index - Child index to insert at.
      * @param structure - Structure of the enum.
      */
-    insertMember(index: number, structure: OptionalKind<EnumMemberStructure>) {
+    insertMember(index: number, structure: OptionalKind<EnumMemberStructure>): EnumMember;
+    /**
+     * Inserts a member to the enum.
+     * @param index - Child index to insert at.
+     * @param structure - Structure of the enum.
+     */
+    insertMember(index: number, structure: OptionalKind<EnumMemberStructure> | WriterFunction): EnumMember | CommentEnumMember;
+    insertMember(index: number, structure: OptionalKind<EnumMemberStructure> | WriterFunction) {
         return this.insertMembers(index, [structure])[0];
     }
 
@@ -63,13 +83,20 @@ export class EnumDeclaration extends EnumDeclarationBase<ts.EnumDeclaration> {
      * @param index - Child index to insert at.
      * @param structures - Structures of the enums.
      */
-    insertMembers(index: number, structures: ReadonlyArray<OptionalKind<EnumMemberStructure>>) {
-        const members = this.getMembers();
-        index = verifyAndGetIndex(index, members.length);
-
+    insertMembers(index: number, structures: ReadonlyArray<OptionalKind<EnumMemberStructure>>): EnumMember[];
+    /**
+     * Inserts members to an enum.
+     * @param index - Child index to insert at.
+     * @param structures - Structures of the enums.
+     */
+    insertMembers(index: number, structures: ReadonlyArray<OptionalKind<EnumMemberStructure> | WriterFunction>): (EnumMember | CommentEnumMember)[];
+    insertMembers(index: number, structures: ReadonlyArray<OptionalKind<EnumMemberStructure> | WriterFunction>) {
         if (structures.length === 0)
             return [];
 
+        const members = this.getMembersWithComments();
+        index = verifyAndGetIndex(index, members.length);
+
         // create member code
         // todo: pass in the StructureToText to the function below
         const writer = this._getWriterWithChildIndentation();
@@ -86,7 +113,7 @@ export class EnumDeclaration extends EnumDeclarationBase<ts.EnumDeclaration> {
         });
 
         // get the members
-        return getNodesToReturn(this.getMembers(), index, structures.length) as EnumMember[];
+        return getNodesToReturn(this.getMembersWithComments(), index, structures.length) as (EnumMember | CommentEnumMember)[];
     }
 
     /**

diff --git a/src/manipulation/helpers/appendCommaToText.ts b/src/manipulation/helpers/appendCommaToText.ts
@@ -0,0 +1,35 @@
+import { ts } from "../../typescript";
+
+const scanner = ts.createScanner(ts.ScriptTarget.Latest, true);
+
+/**
+ * Appends a comma to the text taking into account the various language aspects.
+ */
+export function appendCommaToText(text: string) {
+    const pos = getAppendCommaPos(text, 0);
+    if (pos === -1)
+        return text;
+
+    return text.substring(0, pos) + "," + text.substring(pos);
+}
+
+/**
+ * Gets the position in the text that a comma could be appended.
+ * @param text - Text to search.
+ * @param searchStartPos - Position to start searching.
+ * @returns The position to append. -1 otherwise.
+ */
+export function getAppendCommaPos(text: string, searchStartPos: number) {
+    scanner.setText(text);
+    scanner.setTextPos(searchStartPos);
+
+    if (scanner.scan() === ts.SyntaxKind.EndOfFileToken)
+        return -1;
+
+    while (scanner.scan() !== ts.SyntaxKind.EndOfFileToken) {
+        // just keep scanning...
+    }
+
+    const pos = scanner.getStartPos();
+    return text[pos - 1] === "," ? -1 : pos;
+}
diff --git a/src/manipulation/helpers/index.ts b/src/manipulation/helpers/index.ts
@@ -1,3 +1,4 @@
+export * from "./appendCommaToText";
 export * from "./getEndIndexFromArray";
 export * from "./getInsertPosFromIndex";
 export * from "./getMixinStructureFunctions";

diff --git a/src/manipulation/manipulations/insertion.ts b/src/manipulation/manipulations/insertion.ts
@@ -2,7 +2,7 @@ import { CodeBlockWriter } from "../../codeBlockWriter";
 import { Node, SourceFile } from "../../compiler";
 import { SyntaxKind, ts } from "../../typescript";
 import { TypeGuards, StringUtils } from "../../utils";
-import { getEndPosFromIndex, getInsertPosFromIndex, getRangeWithoutCommentsFromArray, verifyAndGetIndex } from "../helpers";
+import { getEndPosFromIndex, getInsertPosFromIndex, getRangeWithoutCommentsFromArray, verifyAndGetIndex, appendCommaToText } from "../helpers";
 import { NodeHandlerFactory } from "../nodeHandlers";
 import { InsertionTextManipulator } from "../textManipulators";
 import { doManipulation } from "./doManipulation";
@@ -73,9 +73,6 @@ export interface InsertIntoCommaSeparatedNodesOptions {
     surroundWithSpaces?: boolean;
 }
 
-const endsWithComma = /\,\s*$/;
-const startsWithComma = /^\s*\,/;
-
 export function insertIntoCommaSeparatedNodes(opts: InsertIntoCommaSeparatedNodesOptions) {
     const { currentNodes, insertIndex, parent } = opts;
     const nextNode = currentNodes[insertIndex] as Node | undefined;
@@ -137,9 +134,26 @@ export function insertIntoCommaSeparatedNodes(opts: InsertIntoCommaSeparatedNode
     }
 
     function prependCommaAndSeparator() {
-        if (!startsWithComma.test(newText)) {
-            prependSeparator();
-            newText = `,${newText}`;
+        const originalSourceFileText = parent.getSourceFile().getFullText();
+        const previousNodeNextSibling = previousNode!.getNextSibling();
+        let text = "";
+        if (previousNodeNextSibling != null && previousNodeNextSibling.getKind() === SyntaxKind.CommaToken) {
+            appendNodeTrailingCommentRanges(previousNode!);
+            text += ",";
+            appendNodeTrailingCommentRanges(previousNodeNextSibling);
+        }
+        else {
+            text += ",";
+            appendNodeTrailingCommentRanges(previousNode!);
+        }
+
+        prependSeparator();
+        newText = text + newText;
+
+        function appendNodeTrailingCommentRanges(node: Node) {
+            for (const commentRange of node.getTrailingCommentRanges()) {
+                text += originalSourceFileText.substring(node.getEnd(), commentRange.getPos()) + commentRange.getText();
+            }
         }
     }
 
@@ -149,12 +163,8 @@ export function insertIntoCommaSeparatedNodes(opts: InsertIntoCommaSeparatedNode
     }
 
     function appendCommaAndSeparator() {
-        if (!endsWithComma.test(newText)) {
-            newText = StringUtils.insertAtLastNonWhitespace(newText, ",");
-            appendSeparator();
-        }
-        else
-            appendIndentation();
+        newText = appendCommaToText(newText);
+        appendSeparator();
     }
 
     function appendSeparator() {

diff --git a/src/structurePrinters/NodePrinter.ts b/src/structurePrinters/NodePrinter.ts
@@ -9,36 +9,47 @@ export abstract class NodePrinter<TStructure> extends Printer<TStructure> {
         super();
     }
 
+    printTextWithoutTrivia(writer: CodeBlockWriter, structure: TStructure) {
+        this.printTextInternal(writer, structure);
+    }
+
     printText(writer: CodeBlockWriter, structure: TStructure) {
-        const typedStructure: { leadingTrivia: Structure["leadingTrivia"], trailingTrivia: Structure["trailingTrivia"] } = structure as any;
-        const leadingTrivia = typedStructure && typedStructure.leadingTrivia;
-        const trailingTrivia = typedStructure && typedStructure.trailingTrivia;
+        this.printLeadingTrivia(writer, structure);
+        writer.closeComment();
+        this.printTextInternal(writer, structure);
+        this.printTrailingTrivia(writer, structure);
+    }
+
+    protected abstract printTextInternal(writer: CodeBlockWriter, structure: TStructure): void;
+
+    printLeadingTrivia(writer: CodeBlockWriter, structure: TStructure) {
+        const leadingTrivia = (structure as any)["leadingTrivia"] as Structure["leadingTrivia"];
 
         if (leadingTrivia != null) {
-            printTrivia.call(this, leadingTrivia);
+            this.printTrivia(writer, leadingTrivia);
 
             if (writer.isInComment())
                 writer.closeComment();
         }
+    }
 
-        this.printTextInternal(writer, structure);
+    printTrailingTrivia(writer: CodeBlockWriter, structure: TStructure) {
+        const trailingTrivia = (structure as any)["trailingTrivia"] as Structure["trailingTrivia"];
 
         if (trailingTrivia != null)
-            printTrivia.call(this, trailingTrivia);
-
-        function printTrivia(this: NodePrinter<TStructure>, trivia: Structure["leadingTrivia"]) {
-            if (trivia instanceof Array) {
-                for (let i = 0; i < trivia.length; i++) {
-                    this.printTextOrWriterFunc(writer, trivia[i]);
-                    if (i !== trivia.length - 1)
-                        writer.newLineIfLastNot();
-                }
-            }
-            else {
-                this.printTextOrWriterFunc(writer, trivia);
+            this.printTrivia(writer, trailingTrivia);
+    }
+
+    private printTrivia(writer: CodeBlockWriter, trivia: Structure["leadingTrivia"]) {
+        if (trivia instanceof Array) {
+            for (let i = 0; i < trivia.length; i++) {
+                this.printTextOrWriterFunc(writer, trivia[i]);
+                if (i !== trivia.length - 1)
+                    writer.newLineIfLastNot();
             }
         }
+        else {
+            this.printTextOrWriterFunc(writer, trivia);
+        }
     }
-
-    protected abstract printTextInternal(writer: CodeBlockWriter, structure: TStructure): void;
 }
diff --git a/src/structurePrinters/enum/EnumMemberStructurePrinter.ts b/src/structurePrinters/enum/EnumMemberStructurePrinter.ts
@@ -2,15 +2,21 @@ import { CodeBlockWriter } from "../../codeBlockWriter";
 import { EnumMemberStructure, OptionalKind } from "../../structures";
 import { NodePrinter } from "../NodePrinter";
 import { CommaNewLineSeparatedStructuresPrinter } from "../formatting";
+import { WriterFunction } from "../../types";
 
-export class EnumMemberStructurePrinter extends NodePrinter<OptionalKind<EnumMemberStructure>> {
+export class EnumMemberStructurePrinter extends NodePrinter<OptionalKind<EnumMemberStructure> | WriterFunction> {
     private readonly multipleWriter = new CommaNewLineSeparatedStructuresPrinter(this);
 
-    printTexts(writer: CodeBlockWriter, structures: ReadonlyArray<OptionalKind<EnumMemberStructure>> | undefined) {
+    printTexts(writer: CodeBlockWriter, structures: ReadonlyArray<OptionalKind<EnumMemberStructure> | WriterFunction> | undefined) {
         this.multipleWriter.printText(writer, structures);
     }
 
-    protected printTextInternal(writer: CodeBlockWriter, structure: OptionalKind<EnumMemberStructure>) {
+    protected printTextInternal(writer: CodeBlockWriter, structure: OptionalKind<EnumMemberStructure> | WriterFunction) {
+        if (structure instanceof Function) {
+            structure(writer);
+            return;
+        }
+
         this.factory.forJSDoc().printDocs(writer, structure.docs);
         writer.write(structure.name);
         if (typeof structure.value === "string") {