Convert to async function: handle type arguments to then/catch

src/services/codefixes/convertToAsyncFunction.ts

@@ -336,17 +336,17 @@ namespace ts.codefix {
     }
 
     /**
-     * Transforms the 'x' part of `x.then(...)`, or the 'y()' part of `y.catch(...)`, where 'x' and 'y()' are Promises.
+     * Transforms the 'x' part of `x.then(...)`, or the 'y()' part of `y().catch(...)`, where 'x' and 'y()' are Promises.
      */
     function transformPromiseExpressionOfPropertyAccess(node: Expression, transformer: Transformer, prevArgName?: SynthBindingName): readonly Statement[] {
         if (shouldReturn(node, transformer)) {
             return [createReturn(getSynthesizedDeepClone(node))];
         }
 
-        return createVariableOrAssignmentOrExpressionStatement(prevArgName, createAwait(node));
+        return createVariableOrAssignmentOrExpressionStatement(prevArgName, createAwait(node), /*typeAnnotation*/ undefined);
     }
 
-    function createVariableOrAssignmentOrExpressionStatement(variableName: SynthBindingName | undefined, rightHandSide: Expression): readonly Statement[] {
+    function createVariableOrAssignmentOrExpressionStatement(variableName: SynthBindingName | undefined, rightHandSide: Expression, typeAnnotation: TypeNode | undefined): readonly Statement[] {
         if (!variableName || isEmptyBindingName(variableName)) {
             // if there's no argName to assign to, there still might be side effects
             return [createExpressionStatement(rightHandSide)];
@@ -363,11 +363,22 @@ namespace ts.codefix {
                 createVariableDeclarationList([
                     createVariableDeclaration(
                         getSynthesizedDeepClone(getNode(variableName)),
-                        /*type*/ undefined,
+                        typeAnnotation,
                         rightHandSide)],
                     NodeFlags.Const))];
     }
 
+    function maybeAnnotateAndReturn(expressionToReturn: Expression | undefined, typeAnnotation: TypeNode | undefined): readonly Statement[] {
+        if (typeAnnotation && expressionToReturn) {
+            const name = createOptimisticUniqueName("result");
+            return [
+                ...createVariableOrAssignmentOrExpressionStatement(createSynthIdentifier(name), expressionToReturn, typeAnnotation),
+                createReturn(name)
+            ];
+        }
+        return [createReturn(expressionToReturn)];
+    }
+
     // should be kept up to date with isFixablePromiseArgument in suggestionDiagnostics.ts
     function getTransformationBody(func: Expression, prevArgName: SynthBindingName | undefined, argName: SynthBindingName | undefined, parent: CallExpression, transformer: Transformer): readonly Statement[] {
         switch (func.kind) {
@@ -382,7 +393,7 @@ namespace ts.codefix {
 
                 const synthCall = createCall(getSynthesizedDeepClone(func as Identifier), /*typeArguments*/ undefined, isSynthIdentifier(argName) ? [argName.identifier] : []);
                 if (shouldReturn(parent, transformer)) {
-                    return [createReturn(synthCall)];
+                    return maybeAnnotateAndReturn(synthCall, parent.typeArguments?.[0]);
                 }
 
                 const type = transformer.checker.getTypeAtLocation(func);
@@ -392,7 +403,7 @@ namespace ts.codefix {
                     return silentFail();
                 }
                 const returnType = callSignatures[0].getReturnType();
-                const varDeclOrAssignment = createVariableOrAssignmentOrExpressionStatement(prevArgName, createAwait(synthCall));
+                const varDeclOrAssignment = createVariableOrAssignmentOrExpressionStatement(prevArgName, createAwait(synthCall), parent.typeArguments?.[0]);
                 if (prevArgName) {
                     prevArgName.types.push(returnType);
                 }
@@ -409,10 +420,12 @@ namespace ts.codefix {
                     for (const statement of funcBody.statements) {
                         if (isReturnStatement(statement)) {
                             seenReturnStatement = true;
-                        }
-
-                        if (isReturnStatementWithFixablePromiseHandler(statement)) {
-                            refactoredStmts = refactoredStmts.concat(getInnerTransformationBody(transformer, [statement], prevArgName));
+                            if (isReturnStatementWithFixablePromiseHandler(statement)) {
+                                refactoredStmts = refactoredStmts.concat(getInnerTransformationBody(transformer, [statement], prevArgName));
+                            }
+                            else {
+                                refactoredStmts.push(...maybeAnnotateAndReturn(statement.expression, parent.typeArguments?.[0]));
+                            }
                         }
                         else {
                             refactoredStmts.push(statement);
@@ -440,14 +453,14 @@ namespace ts.codefix {
                     const rightHandSide = getSynthesizedDeepClone(funcBody);
                     const possiblyAwaitedRightHandSide = !!transformer.checker.getPromisedTypeOfPromise(returnType) ? createAwait(rightHandSide) : rightHandSide;
                     if (!shouldReturn(parent, transformer)) {
-                        const transformedStatement = createVariableOrAssignmentOrExpressionStatement(prevArgName, possiblyAwaitedRightHandSide);
+                        const transformedStatement = createVariableOrAssignmentOrExpressionStatement(prevArgName, possiblyAwaitedRightHandSide, /*typeAnnotation*/ undefined);
                         if (prevArgName) {
                             prevArgName.types.push(returnType);
                         }
                         return transformedStatement;
                     }
                     else {
-                        return [createReturn(possiblyAwaitedRightHandSide)];
+                        return maybeAnnotateAndReturn(possiblyAwaitedRightHandSide, parent.typeArguments?.[0]);
                     }
                 }
             }

src/services/suggestionDiagnostics.ts

@@ -133,7 +133,7 @@ namespace ts {
         return !!forEachReturnStatement(body, isReturnStatementWithFixablePromiseHandler);
     }
 
-    export function isReturnStatementWithFixablePromiseHandler(node: Node): node is ReturnStatement {
+    export function isReturnStatementWithFixablePromiseHandler(node: Node): node is ReturnStatement & { expression: CallExpression } {
         return isReturnStatement(node) && !!node.expression && isFixablePromiseHandler(node.expression);
     }
 

src/testRunner/unittests/services/convertToAsyncFunction.ts

@@ -255,7 +255,21 @@ interface String { charAt: any; }
 interface Array<T> {}`
     };
 
-    function testConvertToAsyncFunction(caption: string, text: string, baselineFolder: string, includeLib?: boolean, expectFailure = false, onlyProvideAction = false) {
+    type WithSkipAndOnly<T extends any[]> = ((...args: T) => void) & {
+        skip: (...args: T) => void;
+        only: (...args: T) => void;
+    };
+
+    function createTestWrapper<T extends any[]>(fn: (it: Mocha.PendingTestFunction, ...args: T) => void): WithSkipAndOnly<T> {
+        wrapped.skip = (...args: T) => fn(it.skip, ...args);
+        wrapped.only = (...args: T) => fn(it.only, ...args);
+        return wrapped;
+        function wrapped(...args: T) {
+            return fn(it, ...args);
+        }
+    }
+
+    function testConvertToAsyncFunction(it: Mocha.PendingTestFunction, caption: string, text: string, baselineFolder: string, includeLib?: boolean, expectFailure = false, onlyProvideAction = false) {
         const t = extractTest(text);
         const selectionRange = t.ranges.get("selection")!;
         if (!selectionRange) {
@@ -343,7 +357,19 @@ interface Array<T> {}`
         }
     }
 
-    describe("unittests:: services:: convertToAsyncFunctions", () => {
+    const _testConvertToAsyncFunction = createTestWrapper((it, caption: string, text: string) => {
+        testConvertToAsyncFunction(it, caption, text, "convertToAsyncFunction", /*includeLib*/ true);
+    });
+
+    const _testConvertToAsyncFunctionFailed = createTestWrapper((it, caption: string, text: string) => {
+        testConvertToAsyncFunction(it, caption, text, "convertToAsyncFunction", /*includeLib*/ true, /*expectFailure*/ true);
+    });
+
+    const _testConvertToAsyncFunctionFailedSuggestion = createTestWrapper((it, caption: string, text: string) => {
+        testConvertToAsyncFunction(it, caption, text, "convertToAsyncFunction", /*includeLib*/ true, /*expectFailure*/ true, /*onlyProvideAction*/ true);
+    });
+
+    describe("unittests:: services:: convertToAsyncFunction", () => {
         _testConvertToAsyncFunction("convertToAsyncFunction_basic", `
 function [#|f|](): Promise<void>{
     return fetch('https://typescriptlang.org').then(result => { console.log(result) });
@@ -1352,17 +1378,54 @@ function foo() {
     })
 }
 `);
-    });
 
-    function _testConvertToAsyncFunction(caption: string, text: string) {
-        testConvertToAsyncFunction(caption, text, "convertToAsyncFunction", /*includeLib*/ true);
-    }
+        _testConvertToAsyncFunction("convertToAsyncFunction_thenTypeArgument1", `
+type APIResponse<T> = { success: true, data: T } | { success: false };
 
-    function _testConvertToAsyncFunctionFailed(caption: string, text: string) {
-        testConvertToAsyncFunction(caption, text, "convertToAsyncFunction", /*includeLib*/ true, /*expectFailure*/ true);
-    }
+function wrapResponse<T>(response: T): APIResponse<T> {
+    return { success: true, data: response };
+}
 
-    function _testConvertToAsyncFunctionFailedSuggestion(caption: string, text: string) {
-        testConvertToAsyncFunction(caption, text, "convertToAsyncFunction", /*includeLib*/ true, /*expectFailure*/ true, /*onlyProvideAction*/ true);
-    }
+function [#|get|]() {
+    return Promise.resolve(undefined!).then<APIResponse<{ email: string }>>(wrapResponse);
+}
+`);
+
+        _testConvertToAsyncFunction("convertToAsyncFunction_thenTypeArgument2", `
+type APIResponse<T> = { success: true, data: T } | { success: false };
+
+function wrapResponse<T>(response: T): APIResponse<T> {
+    return { success: true, data: response };
+}
+
+function [#|get|]() {
+    return Promise.resolve(undefined!).then<APIResponse<{ email: string }>>(d => wrapResponse(d));
+}
+`);
+
+        _testConvertToAsyncFunction("convertToAsyncFunction_thenTypeArgument3", `
+type APIResponse<T> = { success: true, data: T } | { success: false };
+
+function wrapResponse<T>(response: T): APIResponse<T> {
+    return { success: true, data: response };
+}
+
+function [#|get|]() {
+    return Promise.resolve(undefined!).then<APIResponse<{ email: string }>>(d => {
+        console.log(d);
+        return wrapResponse(d);
+    });
+}
+`);
+
+        _testConvertToAsyncFunction("convertToAsyncFunction_catchTypeArgument1", `
+type APIResponse<T> = { success: true, data: T } | { success: false };
+
+function [#|get|]() {
+    return Promise
+        .resolve<APIResponse<{ email: string }>>({ success: true, data: { email: "" } })
+        .catch<APIResponse<{ email: string }>>(() => ({ success: false }));
+}
+`);
+    });
 }

tests/baselines/reference/convertToAsyncFunction/convertToAsyncFunction_catchTypeArgument1.ts

@@ -0,0 +1,28 @@
+// ==ORIGINAL==
+
+type APIResponse<T> = { success: true, data: T } | { success: false };
+
+function /*[#|*/get/*|]*/() {
+    return Promise
+        .resolve<APIResponse<{ email: string }>>({ success: true, data: { email: "" } })
+        .catch<APIResponse<{ email: string }>>(() => ({ success: false }));
+}
+
+// ==ASYNC FUNCTION::Convert to async function==
+
+type APIResponse<T> = { success: true, data: T } | { success: false };
+
+async function get() {
+    try {
+        return Promise
+            .resolve<APIResponse<{
+                email: string;
+            }>>({ success: true, data: { email: "" } });
+    }
+    catch (e) {
+        const result: APIResponse<{
+            email: string;
+        }> = ({ success: false });
+        return result;
+    }
+}

tests/baselines/reference/convertToAsyncFunction/convertToAsyncFunction_thenTypeArgument1.ts

@@ -0,0 +1,27 @@
+// ==ORIGINAL==
+
+type APIResponse<T> = { success: true, data: T } | { success: false };
+
+function wrapResponse<T>(response: T): APIResponse<T> {
+    return { success: true, data: response };
+}
+
+function /*[#|*/get/*|]*/() {
+    return Promise.resolve(undefined!).then<APIResponse<{ email: string }>>(wrapResponse);
+}
+
+// ==ASYNC FUNCTION::Convert to async function==
+
+type APIResponse<T> = { success: true, data: T } | { success: false };
+
+function wrapResponse<T>(response: T): APIResponse<T> {
+    return { success: true, data: response };
+}
+
+async function get() {
+    const response = await Promise.resolve((undefined!));
+    const result: APIResponse<{
+        email: string;
+    }> = wrapResponse(response);
+    return result;
+}

tests/baselines/reference/convertToAsyncFunction/convertToAsyncFunction_thenTypeArgument2.ts

@@ -0,0 +1,27 @@
+// ==ORIGINAL==
+
+type APIResponse<T> = { success: true, data: T } | { success: false };
+
+function wrapResponse<T>(response: T): APIResponse<T> {
+    return { success: true, data: response };
+}
+
+function /*[#|*/get/*|]*/() {
+    return Promise.resolve(undefined!).then<APIResponse<{ email: string }>>(d => wrapResponse(d));
+}
+
+// ==ASYNC FUNCTION::Convert to async function==
+
+type APIResponse<T> = { success: true, data: T } | { success: false };
+
+function wrapResponse<T>(response: T): APIResponse<T> {
+    return { success: true, data: response };
+}
+
+async function get() {
+    const d = await Promise.resolve((undefined!));
+    const result: APIResponse<{
+        email: string;
+    }> = wrapResponse(d);
+    return result;
+}

tests/baselines/reference/convertToAsyncFunction/convertToAsyncFunction_thenTypeArgument3.ts

@@ -0,0 +1,31 @@
+// ==ORIGINAL==
+
+type APIResponse<T> = { success: true, data: T } | { success: false };
+
+function wrapResponse<T>(response: T): APIResponse<T> {
+    return { success: true, data: response };
+}
+
+function /*[#|*/get/*|]*/() {
+    return Promise.resolve(undefined!).then<APIResponse<{ email: string }>>(d => {
+        console.log(d);
+        return wrapResponse(d);
+    });
+}
+
+// ==ASYNC FUNCTION::Convert to async function==
+
+type APIResponse<T> = { success: true, data: T } | { success: false };
+
+function wrapResponse<T>(response: T): APIResponse<T> {
+    return { success: true, data: response };
+}
+
+async function get() {
+    const d = await Promise.resolve((undefined!));
+    console.log(d);
+    const result: APIResponse<{
+        email: string;
+    }> = wrapResponse(d);
+    return result;
+}