fix: transform should take into account the node changing kinds

BREAKING CHANGE: `transform` now returns a `Node` instead of `this` because the returned node could be the replaced node. Closes #1248
dsherret · Mar 9, 2022 · 75c4a75 · 75c4a75
1 parent c494733
commit 75c4a75
Show file tree

Hide file tree

Showing 6 changed files with 112 additions and 28 deletions.
diff --git a/deno/ts_morph.d.ts b/deno/ts_morph.d.ts
@@ -3727,9 +3727,10 @@ export declare class Node<NodeType extends ts.Node = ts.Node> {
    */
   formatText(settings?: FormatCodeSettings): void;
   /**
-   * Transforms the node using the compiler api nodes and functions (experimental).
+   * Transforms the node using the compiler api nodes and functions and returns
+   * the node that was transformed (experimental).
    *
-   * WARNING: This will forget descendants of transformed nodes.
+   * WARNING: This will forget descendants of transformed nodes and potentially this node.
    * @example Increments all the numeric literals in a source file.
    * ```ts
    * sourceFile.transform(traversal => {
@@ -3748,7 +3749,7 @@ export declare class Node<NodeType extends ts.Node = ts.Node> {
    * });
    * ```
    */
-  transform(visitNode: (traversal: TransformTraversalControl) => ts.Node): this;
+  transform(visitNode: (traversal: TransformTraversalControl) => ts.Node): Node;
   /** Gets the leading comment ranges of the current node. */
   getLeadingCommentRanges(): CommentRange[];
   /** Gets the trailing comment ranges of the current node. */

diff --git a/deno/ts_morph.js b/deno/ts_morph.js
@@ -3669,12 +3669,30 @@ class Node {
         const transformerFactory = context => {
             return rootNode => innerVisit(rootNode, context);
         };
-        ts.transform(compilerNode, [transformerFactory], this._context.compilerOptions.get());
-        replaceSourceFileTextStraight({
-            sourceFile: this._sourceFile,
-            newText: getTransformedText(),
-        });
-        return this;
+        if (this.getKind() === ts.SyntaxKind.SourceFile) {
+            ts.transform(compilerNode, [transformerFactory], this._context.compilerOptions.get());
+            replaceSourceFileTextStraight({
+                sourceFile: this._sourceFile,
+                newText: getTransformedText([0, this.getEnd()]),
+            });
+            return this;
+        }
+        else {
+            const parent = this.getParentSyntaxList() || this.getParentOrThrow();
+            const childIndex = this.getChildIndex();
+            const start = this.getStart(true);
+            const end = this.getEnd();
+            ts.transform(compilerNode, [transformerFactory], this._context.compilerOptions.get());
+            insertIntoParentTextRange({
+                parent,
+                insertPos: start,
+                newText: getTransformedText([start, end]),
+                replacing: {
+                    textLength: end - start,
+                },
+            });
+            return parent.getChildren()[childIndex];
+        }
         function innerVisit(node, context) {
             const traversal = {
                 visitChildren() {
@@ -3708,16 +3726,16 @@ class Node {
                     wrappedNode.forgetDescendants();
             }
         }
-        function getTransformedText() {
+        function getTransformedText(replaceRange) {
             const fileText = compilerSourceFile.getFullText();
             let finalText = "";
-            let lastPos = 0;
+            let lastPos = replaceRange[0];
             for (const transform of transformations) {
                 finalText += fileText.substring(lastPos, transform.start);
                 finalText += printer.printNode(ts.EmitHint.Unspecified, transform.compilerNode, compilerSourceFile);
                 lastPos = transform.end;
             }
-            finalText += fileText.substring(lastPos);
+            finalText += fileText.substring(lastPos, replaceRange[1]);
             return finalText;
         }
     }

diff --git a/packages/scripts/deps.ts b/packages/scripts/deps.ts
@@ -1,2 +1,2 @@
 export * as path from "https://deno.land/std@0.128.0/path/mod.ts";
-export * as tsMorph from "https://deno.land/x/ts_morph@13.0.2/mod.ts";
+export * as tsMorph from "https://deno.land/x/ts_morph@14.0.0/mod.ts";
diff --git a/packages/ts-morph/lib/ts-morph.d.ts b/packages/ts-morph/lib/ts-morph.d.ts
@@ -3727,9 +3727,10 @@ export declare class Node<NodeType extends ts.Node = ts.Node> {
    */
   formatText(settings?: FormatCodeSettings): void;
   /**
-   * Transforms the node using the compiler api nodes and functions (experimental).
+   * Transforms the node using the compiler api nodes and functions and returns
+   * the node that was transformed (experimental).
    *
-   * WARNING: This will forget descendants of transformed nodes.
+   * WARNING: This will forget descendants of transformed nodes and potentially this node.
    * @example Increments all the numeric literals in a source file.
    * ```ts
    * sourceFile.transform(traversal => {
@@ -3748,7 +3749,7 @@ export declare class Node<NodeType extends ts.Node = ts.Node> {
    * });
    * ```
    */
-  transform(visitNode: (traversal: TransformTraversalControl) => ts.Node): this;
+  transform(visitNode: (traversal: TransformTraversalControl) => ts.Node): Node;
   /** Gets the leading comment ranges of the current node. */
   getLeadingCommentRanges(): CommentRange[];
   /** Gets the trailing comment ranges of the current node. */

diff --git a/packages/ts-morph/src/compiler/ast/common/Node.ts b/packages/ts-morph/src/compiler/ast/common/Node.ts
@@ -1439,9 +1439,10 @@ export class Node<NodeType extends ts.Node = ts.Node> {
     }
 
     /**
-     * Transforms the node using the compiler api nodes and functions (experimental).
+     * Transforms the node using the compiler api nodes and functions and returns
+     * the node that was transformed (experimental).
      *
-     * WARNING: This will forget descendants of transformed nodes.
+     * WARNING: This will forget descendants of transformed nodes and potentially this node.
      * @example Increments all the numeric literals in a source file.
      * ```ts
      * sourceFile.transform(traversal => {
@@ -1460,7 +1461,7 @@ export class Node<NodeType extends ts.Node = ts.Node> {
      * });
      * ```
      */
-    transform(visitNode: (traversal: TransformTraversalControl) => ts.Node) {
+    transform(visitNode: (traversal: TransformTraversalControl) => ts.Node): Node {
         const compilerFactory = this._context.compilerFactory;
         const printer = ts.createPrinter({
             newLine: this._context.manipulationSettings.getNewLineKind(),
@@ -1474,14 +1475,34 @@ export class Node<NodeType extends ts.Node = ts.Node> {
             return rootNode => innerVisit(rootNode, context);
         };
 
-        ts.transform(compilerNode, [transformerFactory], this._context.compilerOptions.get());
+        if (this.getKind() === ts.SyntaxKind.SourceFile) {
+          ts.transform(compilerNode, [transformerFactory], this._context.compilerOptions.get());
 
-        replaceSourceFileTextStraight({
-            sourceFile: this._sourceFile,
-            newText: getTransformedText(),
-        });
+          replaceSourceFileTextStraight({
+              sourceFile: this._sourceFile,
+              newText: getTransformedText([0, this.getEnd()]),
+          });
 
-        return this;
+          return this;
+        } else {
+          const parent = this.getParentSyntaxList() || this.getParentOrThrow();
+          const childIndex = this.getChildIndex();
+          const start = this.getStart(true);
+          const end = this.getEnd();
+
+          ts.transform(compilerNode, [transformerFactory], this._context.compilerOptions.get());
+
+          insertIntoParentTextRange({
+              parent,
+              insertPos: start,
+              newText: getTransformedText([start, end]),
+              replacing: {
+                  textLength: end - start,
+              },
+          });
+
+          return parent.getChildren()[childIndex];
+        }
 
         function innerVisit(node: ts.Node, context: ts.TransformationContext) {
             const traversal: TransformTraversalControl = {
@@ -1525,18 +1546,18 @@ export class Node<NodeType extends ts.Node = ts.Node> {
             }
         }
 
-        function getTransformedText() {
+        function getTransformedText(replaceRange: [number, number]) {
             const fileText = compilerSourceFile.getFullText();
             let finalText = "";
-            let lastPos = 0;
+            let lastPos = replaceRange[0];
 
             for (const transform of transformations) {
                 finalText += fileText.substring(lastPos, transform.start);
                 finalText += printer.printNode(ts.EmitHint.Unspecified, transform.compilerNode, compilerSourceFile);
                 lastPos = transform.end;
             }
 
-            finalText += fileText.substring(lastPos);
+            finalText += fileText.substring(lastPos, replaceRange[1]);
             return finalText;
         }
     }

diff --git a/packages/ts-morph/src/tests/compiler/ast/common/nodeTests.ts b/packages/ts-morph/src/tests/compiler/ast/common/nodeTests.ts
@@ -2171,6 +2171,49 @@ class MyClass {
         return node;
       }, "hello");
     });
+
+    it("should handle changing the node", () => {
+      const { sourceFile } = getInfoFromText("export class C {}");
+      const node = sourceFile.getClassOrThrow("C");
+
+      const newNode = node.transform(() => {
+        return ts.createFunctionDeclaration(
+          undefined,
+          undefined,
+          undefined,
+          "test",
+          undefined,
+          [],
+          undefined,
+          ts.createBlock([]),
+        );
+      });
+      expect(newNode.getText()).to.equal("function test() { }");
+      expect(sourceFile.getText()).to.equal("function test() { }");
+      expect(node.wasForgotten()).to.be.true;
+    });
+
+    it("should remember the node when not changing kind", () => {
+      const { sourceFile } = getInfoFromText("function original() {}");
+      const node = sourceFile.getFunctionOrThrow("original");
+
+      const newNode = node.transform(() => {
+        return ts.createFunctionDeclaration(
+          undefined,
+          undefined,
+          undefined,
+          "test",
+          undefined,
+          [],
+          undefined,
+          ts.createBlock([]),
+        );
+      });
+      expect(newNode.getText()).to.equal("function test() { }");
+      expect(sourceFile.getText()).to.equal("function test() { }");
+      expect(node.wasForgotten()).to.be.false;
+      expect(node === newNode).to.be.true;
+    });
   });
 
   describe(nameof<Node>("getLocals"), () => {