fix: properly insert async keyword in arrow function (#1597)

dsherret · Dec 30, 2024 · 84613e5 · 84613e5
1 parent 04a1679
commit 84613e5
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 4 deletions.
diff --git a/packages/ts-morph/src/compiler/ast/base/ModifierableNode.ts b/packages/ts-morph/src/compiler/ast/base/ModifierableNode.ts
@@ -161,6 +161,8 @@ export function ModifierableNode<T extends Constructor<ModifierableNodeExtension
         function getInitialInsertPos() {
           if (modifiers.length > 0)
             return modifiers[0].getStart();
+          if (node.getKind() === SyntaxKind.ArrowFunction)
+            return node.getStart();
           for (const child of node._getChildrenIterator()) {
             // skip over any initial syntax lists (ex. decorators) or js docs
             if (child.getKind() === SyntaxKind.SyntaxList || ts.isJSDocCommentContainingNode(child.compilerNode))

diff --git a/packages/ts-morph/src/tests/compiler/ast/base/asyncableNodeTests.ts b/packages/ts-morph/src/tests/compiler/ast/base/asyncableNodeTests.ts
@@ -1,6 +1,6 @@
-import { nameof } from "@ts-morph/common";
+import { nameof, SyntaxKind } from "@ts-morph/common";
 import { expect } from "chai";
-import { AsyncableNode, FunctionDeclaration } from "../../../../compiler";
+import { AsyncableNode, FunctionDeclaration, Node, VariableStatement } from "../../../../compiler";
 import { AsyncableNodeStructure } from "../../../../structures";
 import { getInfoFromText } from "../../testHelpers";
 
@@ -43,8 +43,10 @@ describe("AsyncableNode", () => {
 
   describe(nameof<AsyncableNode>("setIsAsync"), () => {
     function doTest(text: string, value: boolean, expected: string) {
-      const { firstChild, sourceFile } = getInfoFromText<FunctionDeclaration>(text);
-      firstChild.setIsAsync(value);
+      let { firstChild, sourceFile } = getInfoFromText(text);
+      if (firstChild.getKind() === SyntaxKind.VariableStatement)
+        firstChild = (firstChild as VariableStatement).getDeclarations()[0].getInitializerOrThrow();
+      (firstChild as any as AsyncableNode).setIsAsync(value);
       expect(sourceFile.getText()).to.equal(expected);
     }
 
@@ -55,6 +57,12 @@ describe("AsyncableNode", () => {
     it("should set as not async when async", () => {
       doTest("async function Identifier() {}", false, "function Identifier() {}");
     });
+
+    it("should handle arrow function", () => {
+      doTest("const f = a => a * 2;", true, "const f = async a => a * 2;");
+      doTest("const f = async a => a * 2;", true, "const f = async a => a * 2;");
+      doTest("const f = async a => a * 2;", false, "const f = a => a * 2;");
+    });
   });
 
   describe(nameof<FunctionDeclaration>("getStructure"), () => {