Change explicit return type T to Promise<T>

Author: srolel

src/services/codefixes/fixAwaitInSyncFunction.ts

@@ -9,19 +9,42 @@ namespace ts.codefix {
         errorCodes,
         getCodeActions(context) {
             const { sourceFile, span } = context;
-            const node = getNodeToInsertBefore(sourceFile, span.start);
-            if (!node) return undefined;
-            const changes = textChanges.ChangeTracker.with(context, t => doChange(t, sourceFile, node));
+            const token = getTokenAtPosition(sourceFile, span.start, /*includeJsDocComment*/ false);
+            const containingFunction = getContainingFunction(token);
+            const insertBefore = getNodeToInsertBefore(sourceFile, containingFunction);
+            const returnType = getReturnTypeNode(containingFunction);
+            if (!insertBefore) return undefined;
+            const changes = textChanges.ChangeTracker.with(context, t => doChange(t, sourceFile, insertBefore, returnType));
             return [{ description: getLocaleSpecificMessage(Diagnostics.Convert_to_async), changes, fixId }];
         },
         fixIds: [fixId],
-        getAllCodeActions: context => codeFixAll(context, errorCodes, (changes, diag) =>
-            doChange(changes, context.sourceFile, getNodeToInsertBefore(diag.file, diag.start!))),
+        getAllCodeActions: context => codeFixAll(context, errorCodes, (changes, diag) => {
+            const token = getTokenAtPosition(diag.file, diag.start!, /*includeJsDocComment*/ false);
+            const containingFunction = getContainingFunction(token);
+            const insertBefore = getNodeToInsertBefore(diag.file, containingFunction);
+            const returnType = getReturnTypeNode(containingFunction);
+            if (insertBefore) {
+                doChange(changes, context.sourceFile, insertBefore, returnType);
+            }
+        }),
     });
 
-    function getNodeToInsertBefore(sourceFile: SourceFile, pos: number): Node | undefined {
-        const token = getTokenAtPosition(sourceFile, pos, /*includeJsDocComment*/ false);
-        const containingFunction = getContainingFunction(token);
+    function getReturnTypeNode(containingFunction: FunctionLike): TypeNode | undefined {
+        switch (containingFunction.kind) {
+            case SyntaxKind.MethodDeclaration:
+            case SyntaxKind.FunctionDeclaration:
+                return containingFunction.type;
+            case SyntaxKind.FunctionExpression:
+            case SyntaxKind.ArrowFunction:
+                if (isVariableDeclaration(containingFunction.parent) &&
+                    containingFunction.parent.type &&
+                    isFunctionTypeNode(containingFunction.parent.type)) {
+                    return containingFunction.parent.type.type;
+                }
+        }
+    }
+
+    function getNodeToInsertBefore(sourceFile: SourceFile, containingFunction: FunctionLike): Node | undefined {
         switch (containingFunction.kind) {
             case SyntaxKind.MethodDeclaration:
                 return containingFunction.name;
@@ -35,7 +58,13 @@ namespace ts.codefix {
         }
     }
 
-    function doChange(changes: textChanges.ChangeTracker, sourceFile: SourceFile, insertBefore: Node): void {
+    function doChange(changes: textChanges.ChangeTracker, sourceFile: SourceFile, insertBefore: Node, returnType: TypeNode | undefined): void {
+        if (returnType) {
+            const entityName = getEntityNameFromTypeNode(returnType);
+            if (!entityName || entityName.getText() !== "Promise") {
+                changes.replaceNode(sourceFile, returnType, createTypeReferenceNode("Promise", createNodeArray([returnType])));
+            }
+        }
         changes.insertModifierBefore(sourceFile, SyntaxKind.AsyncKeyword, insertBefore);
     }
 }

tests/cases/fourslash/codeFixAwaitInSyncFunction10.ts

@@ -0,0 +1,13 @@
+/// <reference path='fourslash.ts' />
+
+////const f: () => number | string = () => {
+////    await Promise.resolve('foo');
+////}
+
+verify.codeFix({
+    description: "Convert to async",
+    newFileContent:
+`const f: () => Promise<number | string> = async () => {
+    await Promise.resolve('foo');
+}`,
+});

tests/cases/fourslash/codeFixAwaitInSyncFunction11.ts

@@ -0,0 +1,14 @@
+/// <reference path='fourslash.ts' />
+
+////const f: string = () => {
+////    await Promise.resolve('foo');
+////}
+
+// should not change type if it's incorrectly set
+verify.codeFix({
+    description: "Convert to async",
+    newFileContent:
+`const f: string = async () => {
+    await Promise.resolve('foo');
+}`,
+});

tests/cases/fourslash/codeFixAwaitInSyncFunction12.ts

@@ -0,0 +1,13 @@
+/// <reference path='fourslash.ts' />
+
+////const f: () => Array<number | string> = function() {
+////    await Promise.resolve([]);
+////}
+
+verify.codeFix({
+    description: "Convert to async",
+    newFileContent:
+`const f: () => Promise<Array<number | string>> = async function() {
+    await Promise.resolve([]);
+}`,
+});

tests/cases/fourslash/codeFixAwaitInSyncFunction13.ts

@@ -0,0 +1,13 @@
+/// <reference path='fourslash.ts' />
+
+////const f: () => Promise<number | string> = () => {
+////    await Promise.resolve('foo');
+////}
+
+verify.codeFix({
+    description: "Convert to async",
+    newFileContent:
+`const f: () => Promise<number | string> = async () => {
+    await Promise.resolve('foo');
+}`,
+});

tests/cases/fourslash/codeFixAwaitInSyncFunction8.ts

@@ -0,0 +1,13 @@
+/// <reference path='fourslash.ts' />
+
+////function f(): number | string {
+////    await Promise.resolve(8);
+////}
+
+verify.codeFix({
+    description: "Convert to async",
+    newFileContent:
+`async function f(): Promise<number | string> {
+    await Promise.resolve(8);
+}`,
+});

tests/cases/fourslash/codeFixAwaitInSyncFunction9.ts

@@ -0,0 +1,17 @@
+/// <reference path='fourslash.ts' />
+
+////class Foo {
+////    bar(): string {
+////        await Promise.resolve('baz');
+////    }
+////}
+
+verify.codeFix({
+    description: "Convert to async",
+    newFileContent:
+`class Foo {
+    async bar(): Promise<string> {
+        await Promise.resolve('baz');
+    }
+}`,
+});