fix: support type parameters on arrow functions

packages/ts-morph/src/compiler/ast/base/TypeParameteredNode.ts

@@ -104,6 +104,11 @@ export function TypeParameteredNode<T extends Constructor<TypeParameteredNodeExt
       structurePrinter.printTexts(writer, structures);
 
       if (typeParameters.length === 0) {
+        if (Node.isArrowFunction(this)) {
+          addArrowFunctionParensIfNecessary(this);
+          if (this._sourceFile.getLanguageVariant() === ts.LanguageVariant.JSX && structures.length === 1)
+            writer.write(",");
+        }
         insertIntoParentTextRange({
           insertPos: getInsertPos(this),
           parent: this,
@@ -145,8 +150,30 @@ function getInsertPos(node: TypeParameteredNode & Node) {
   const namedNode = node as any as (NamedNode & Node);
   if (namedNode.getNameNode != null)
     return namedNode.getNameNode().getEnd();
-  else if (Node.isCallSignatureDeclaration(node) || Node.isFunctionTypeNode(node))
+  else if (Node.isCallSignatureDeclaration(node) || Node.isFunctionTypeNode(node) || Node.isArrowFunction(node))
     return node.getFirstChildByKindOrThrow(SyntaxKind.OpenParenToken).getStart();
   else
     throw new errors.NotImplementedError(`Not implemented scenario inserting type parameters for node with kind ${node.getKindName()}.`);
 }
+
+function addArrowFunctionParensIfNecessary(node: Node<ts.ArrowFunction>) {
+  if (node.getFirstChildByKind(SyntaxKind.OpenParenToken) != null)
+    return;
+
+  const parameter = node._getNodeFromCompilerNode(node.compilerNode.parameters[0]);
+  const name = parameter.getNameNode();
+  const text = parameter.getText();
+  insertIntoParentTextRange({
+    parent: node,
+    insertPos: parameter.getStart(),
+    newText: `(${text})`,
+    replacing: { textLength: text.length },
+    customMappings: newParent => {
+      const newParameter = (newParent as ts.ArrowFunction).parameters[0];
+      return [
+        { currentNode: parameter, newNode: newParameter },
+        { currentNode: name, newNode: newParameter.name },
+      ];
+    },
+  });
+}

packages/ts-morph/src/tests/compiler/ast/function/arrowFunctionTests.ts

@@ -9,6 +9,75 @@ function getInfoFromTextWithExpression(text: string) {
 }
 
 describe("ArrowFunction", () => {
+  describe(nameof<ArrowFunction>("addTypeParameters"), () => {
+    for (
+      const [text, expectedText] of [
+        ["(value) => value", "<T>(value) => value"],
+        ["() => 1", "<T>() => 1"],
+        ["async (value) => value", "async <T>(value) => value"],
+        ["value => value", "<T>(value) => value"],
+        ["async value => value", "async <T>(value) => value"],
+        ["value /* comment */ => value", "<T>(value) /* comment */ => value"],
+      ]
+    ) {
+      it(`should add type parameters to ${text}`, () => {
+        const { expression } = getInfoFromTextWithExpression(text);
+        const parameter = expression.getParameters()[0];
+        const name = parameter?.getNameNode();
+        const result = expression.addTypeParameters(["T"]);
+        expect(expression.getText()).to.equal(expectedText);
+        expect(result.map(p => p.getName())).to.deep.equal(["T"]);
+        expect(expression.getTypeParameters()[0]).to.equal(result[0]);
+        if (parameter != null) {
+          expect(expression.getParameters()[0]).to.equal(parameter);
+          expect(parameter.getNameNode()).to.equal(name);
+          expect(parameter.getText()).to.equal("value");
+        }
+      });
+    }
+
+    for (const prefix of ["", "async "]) {
+      it(`should preserve a cached parameter name in ${prefix}value => value`, () => {
+        const { expression } = getInfoFromTextWithExpression(`${prefix}value => value`);
+        const name = expression.getParameters()[0].getNameNode().asKindOrThrow(SyntaxKind.Identifier);
+        expression.addTypeParameters(["T"]);
+        name.rename("item");
+        expect(expression.getText()).to.equal(`${prefix}<T>(item) => item`);
+      });
+    }
+
+    it("should disambiguate a single type parameter from JSX", () => {
+      const { descendant } = getInfoFromTextWithDescendant<ArrowFunction>("value => value", SyntaxKind.ArrowFunction, { isJsx: true });
+      descendant.addTypeParameters(["T"]);
+      expect(descendant.getText()).to.equal("<T,>(value) => value");
+      expect(descendant.getTypeParameters()[0].getName()).to.equal("T");
+    });
+
+    it("should add multiple type parameters in JSX", () => {
+      const { descendant } = getInfoFromTextWithDescendant<ArrowFunction>("(value) => value", SyntaxKind.ArrowFunction, { isJsx: true });
+      descendant.addTypeParameters(["T", "U extends T"]);
+      expect(descendant.getText()).to.equal("<T, U extends T>(value) => value");
+    });
+
+    it("should append to existing type parameters", () => {
+      const { expression } = getInfoFromTextWithExpression("<T>(value: T) => value");
+      expression.addTypeParameters(["U extends T"]);
+      expect(expression.getText()).to.equal("<T, U extends T>(value: T) => value");
+    });
+
+    it("should leave an unparenthesized parameter unchanged when no type parameters are added", () => {
+      const { expression } = getInfoFromTextWithExpression("value => value");
+      expect(expression.addTypeParameters([])).to.deep.equal([]);
+      expect(expression.getText()).to.equal("value => value");
+    });
+
+    it("should leave an unparenthesized parameter unchanged when the index is invalid", () => {
+      const { expression } = getInfoFromTextWithExpression("value => value");
+      expect(() => expression.insertTypeParameter(1, "T")).to.throw();
+      expect(expression.getText()).to.equal("value => value");
+    });
+  });
+
   describe(nameof<ArrowFunction>("getEqualsGreaterThan"), () => {
     function doTest(text: string, expectedText: string) {
       const { expression } = getInfoFromTextWithExpression(text);