fix(42519): fix infer function return type quick fix with arrow functions without parens

Author: a-tarasyuk

src/services/refactors/inferFunctionReturnType.ts

@@ -17,8 +17,7 @@ namespace ts.refactor.inferFunctionReturnType {
     function getEditsForAction(context: RefactorContext): RefactorEditInfo | undefined {
         const info = getInfo(context);
         if (info && !isRefactorErrorInfo(info)) {
-            const edits = textChanges.ChangeTracker.with(context, t =>
-                t.tryInsertTypeAnnotation(context.file, info.declaration, info.returnTypeNode));
+            const edits = textChanges.ChangeTracker.with(context, t => doChange(context.file, t, info.declaration, info.returnTypeNode));
             return { renameFilename: undefined, renameLocation: undefined, edits };
         }
         return undefined;
@@ -55,6 +54,19 @@ namespace ts.refactor.inferFunctionReturnType {
         returnTypeNode: TypeNode;
     }
 
+    function doChange(sourceFile: SourceFile, changes: textChanges.ChangeTracker, declaration: ConvertibleDeclaration, typeNode: TypeNode) {
+        const closeParen = findChildOfKind(declaration, SyntaxKind.CloseParenToken, sourceFile);
+        const needParens = isArrowFunction(declaration) && closeParen === undefined;
+        const endNode = needParens ? first(declaration.parameters) : closeParen;
+        if (endNode) {
+            if (needParens) {
+                changes.insertNodeBefore(sourceFile, endNode, factory.createToken(SyntaxKind.OpenParenToken));
+                changes.insertNodeAfter(sourceFile, endNode, factory.createToken(SyntaxKind.CloseParenToken));
+            }
+            changes.insertNodeAt(sourceFile, endNode.end, typeNode, { prefix: ": " });
+        }
+    }
+
     function getInfo(context: RefactorContext): FunctionInfo | RefactorErrorInfo | undefined {
         if (isInJSFile(context.file) || !refactorKindBeginsWith(inferReturnTypeAction.kind, context.kind)) return;
 

tests/cases/fourslash/refactorInferFunctionReturnType22.ts

@@ -0,0 +1,16 @@
+/// <reference path='fourslash.ts' />
+
+////const foo = async /*a*//*b*/a => {
+////    return 1;
+////}
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Infer function return type",
+    actionName: "Infer function return type",
+    actionDescription: "Infer function return type",
+    newContent:
+`const foo = async (a): Promise<number> => {
+    return 1;
+}`
+});

tests/cases/fourslash/refactorInferFunctionReturnType23.ts

@@ -0,0 +1,16 @@
+/// <reference path='fourslash.ts' />
+
+////const foo = async /*a*//*b*/a => {
+////    return 1;
+////}
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Infer function return type",
+    actionName: "Infer function return type",
+    actionDescription: "Infer function return type",
+    newContent:
+`const foo = async (a): Promise<number> => {
+    return 1;
+}`
+});

tests/cases/fourslash/refactorInferFunctionReturnType24.ts

@@ -0,0 +1,16 @@
+/// <reference path='fourslash.ts' />
+
+////const foo = async /*a*//*b*/(a) => {
+////    return 1;
+////}
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Infer function return type",
+    actionName: "Infer function return type",
+    actionDescription: "Infer function return type",
+    newContent:
+`const foo = async (a): Promise<number> => {
+    return 1;
+}`
+});

tests/cases/fourslash/refactorInferFunctionReturnType25.ts

@@ -0,0 +1,16 @@
+/// <reference path='fourslash.ts' />
+
+////const foo = async /*a*//*b*/(a) => {
+////    return 1;
+////}
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Infer function return type",
+    actionName: "Infer function return type",
+    actionDescription: "Infer function return type",
+    newContent:
+`const foo = async (a): Promise<number> => {
+    return 1;
+}`
+});