add fixAwaitInSyncFunction code fix

src/compiler/diagnosticMessages.json

@@ -3843,6 +3843,10 @@
         "category": "Message",
         "code": 90028
     },
+    "Add async modifier to containing function": {
+        "category": "Message",
+        "code": 90029
+    },
     "Convert function to an ES2015 class": {
         "category": "Message",
         "code": 95001

src/services/codefixes/fixAwaitInSyncFunction.ts

@@ -0,0 +1,74 @@
+/* @internal */
+namespace ts.codefix {
+    const fixId = "fixAwaitInSyncFunction";
+    const errorCodes = [
+        Diagnostics.await_expression_is_only_allowed_within_an_async_function.code,
+        Diagnostics.A_for_await_of_statement_is_only_allowed_within_an_async_function_or_async_generator.code,
+    ];
+    registerCodeFix({
+        errorCodes,
+        getCodeActions(context) {
+            const { sourceFile, span } = context;
+            const nodes = getNodes(sourceFile, span.start);
+            if (!nodes) return undefined;
+            const changes = textChanges.ChangeTracker.with(context, t => doChange(t, sourceFile, nodes));
+            return [{ description: getLocaleSpecificMessage(Diagnostics.Add_async_modifier_to_containing_function), changes, fixId }];
+        },
+        fixIds: [fixId],
+        getAllCodeActions: context => codeFixAll(context, errorCodes, (changes, diag) => {
+            const nodes = getNodes(diag.file, diag.start);
+            if (!nodes) return;
+            doChange(changes, context.sourceFile, nodes);
+        }),
+    });
+
+    function getReturnType(expr: FunctionDeclaration | MethodDeclaration | FunctionExpression | ArrowFunction) {
+        if (expr.type) {
+            return expr.type;
+        }
+        if (isVariableDeclaration(expr.parent) &&
+            expr.parent.type &&
+            isFunctionTypeNode(expr.parent.type)) {
+            return expr.parent.type.type;
+        }
+    }
+
+    function getNodes(sourceFile: SourceFile, start: number): { insertBefore: Node, returnType: TypeNode | undefined } | undefined {
+        const token = getTokenAtPosition(sourceFile, start, /*includeJsDocComment*/ false);
+        const containingFunction = getContainingFunction(token);
+        let insertBefore: Node | undefined;
+        switch (containingFunction.kind) {
+            case SyntaxKind.MethodDeclaration:
+                insertBefore = containingFunction.name;
+                break;
+            case SyntaxKind.FunctionDeclaration:
+            case SyntaxKind.FunctionExpression:
+                insertBefore = findChildOfKind(containingFunction, SyntaxKind.FunctionKeyword, sourceFile);
+                break;
+            case SyntaxKind.ArrowFunction:
+                insertBefore = findChildOfKind(containingFunction, SyntaxKind.OpenParenToken, sourceFile) || first(containingFunction.parameters);
+                break;
+            default:
+                return;
+        }
+
+        return {
+            insertBefore,
+            returnType: getReturnType(containingFunction)
+        };
+    }
+
+    function doChange(
+        changes: textChanges.ChangeTracker,
+        sourceFile: SourceFile,
+        { insertBefore, returnType }: { insertBefore: Node | undefined, returnType: TypeNode | undefined }): void {
+
+        if (returnType) {
+            const entityName = getEntityNameFromTypeNode(returnType);
+            if (!entityName || entityName.kind !== SyntaxKind.Identifier || entityName.text !== "Promise") {
+                changes.replaceNode(sourceFile, returnType, createTypeReferenceNode("Promise", createNodeArray([returnType])));
+            }
+        }
+        changes.insertModifierBefore(sourceFile, SyntaxKind.AsyncKeyword, insertBefore);
+    }
+}

src/services/codefixes/fixes.ts

@@ -11,6 +11,7 @@
 /// <reference path="fixForgottenThisPropertyAccess.ts" />
 /// <reference path='fixUnusedIdentifier.ts' />
 /// <reference path='fixJSDocTypes.ts' />
+/// <reference path='fixAwaitInSyncFunction.ts' />
 /// <reference path='importFixes.ts' />
 /// <reference path='disableJsDiagnostics.ts' />
 /// <reference path='helpers.ts' />

src/services/textChanges.ts

@@ -345,6 +345,11 @@ namespace ts.textChanges {
             return this.replaceWithSingle(sourceFile, startPosition, startPosition, newNode, this.getOptionsForInsertNodeBefore(before, blankLineBetween));
         }
 
+        public insertModifierBefore(sourceFile: SourceFile, modifier: SyntaxKind, before: Node): void {
+            const pos = before.getStart(sourceFile);
+            this.replaceWithSingle(sourceFile, pos, pos, createToken(modifier), { suffix: " " });
+        }
+
         public changeIdentifierToPropertyAccess(sourceFile: SourceFile, prefix: string, node: Identifier): void {
             const startPosition = getAdjustedStartPosition(sourceFile, node, {}, Position.Start);
             this.replaceWithSingle(sourceFile, startPosition, startPosition, createPropertyAccess(createIdentifier(prefix), ""), {});

tests/cases/fourslash/codeFixAwaitInSyncFunction1.ts

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

tests/cases/fourslash/codeFixAwaitInSyncFunction10.ts

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

tests/cases/fourslash/codeFixAwaitInSyncFunction14.ts

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

tests/cases/fourslash/codeFixAwaitInSyncFunction15.ts

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

tests/cases/fourslash/codeFixAwaitInSyncFunction2.ts

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

tests/cases/fourslash/codeFixAwaitInSyncFunction3.ts

@@ -0,0 +1,12 @@
+/// <reference path='fourslash.ts' />
+
+////const f = {
+////    get a() {
+////        return await Promise.resolve();
+////    },
+////    get a() {
+////        await Promise.resolve();
+////    },
+////}
+
+verify.not.codeFixAvailable();

tests/cases/fourslash/codeFixAwaitInSyncFunction4.ts

@@ -0,0 +1,9 @@
+/// <reference path='fourslash.ts' />
+
+////class Foo {
+////    constructor {
+////        await Promise.resolve();
+////    }
+////}
+
+verify.not.codeFixAvailable();

tests/cases/fourslash/codeFixAwaitInSyncFunction5.ts

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

tests/cases/fourslash/codeFixAwaitInSyncFunction6.5.ts

@@ -0,0 +1,13 @@
+/// <reference path='fourslash.ts' />
+
+////const f = promise => {
+////    await promise;
+////}
+
+verify.codeFix({
+    description: "Add async modifier to containing function",
+    newFileContent:
+`const f = async promise => {
+    await promise;
+}`,
+});

tests/cases/fourslash/codeFixAwaitInSyncFunction6.ts

@@ -0,0 +1,13 @@
+/// <reference path='fourslash.ts' />
+
+////const f = (promise) => {
+////    await promise;
+////}
+
+verify.codeFix({
+    description: "Add async modifier to containing function",
+    newFileContent:
+`const f = async (promise) => {
+    await promise;
+}`,
+});

tests/cases/fourslash/codeFixAwaitInSyncFunction7.ts

@@ -0,0 +1,17 @@
+/// <reference path='fourslash.ts' />
+
+////function f() {
+////    for await (const x of g()) {
+////        console.log(x);
+////    }
+////}
+
+verify.codeFix({
+    description: "Add async modifier to containing function",
+    newFileContent:
+`async function f() {
+    for await (const x of g()) {
+        console.log(x);
+    }
+}`,
+});

tests/cases/fourslash/codeFixAwaitInSyncFunction8.ts

@@ -0,0 +1,13 @@
+/// <reference path='fourslash.ts' />
+
+////function f(): number | string {
+////    await Promise.resolve(8);
+////}
+
+verify.codeFix({
+    description: "Add async modifier to containing function",
+    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: "Add async modifier to containing function",
+    newFileContent:
+`class Foo {
+    async bar(): Promise<string> {
+        await Promise.resolve('baz');
+    }
+}`,
+});

tests/cases/fourslash/codeFixAwaitInSyncFunction_all.ts

@@ -0,0 +1,21 @@
+/// <reference path='fourslash.ts' />
+
+////function f() {
+////    await Promise.resolve();
+////}
+////
+////const g = () => {
+////    await f();
+////}
+
+verify.codeFixAll({
+    fixId: "fixAwaitInSyncFunction",
+    newFileContent:
+`async function f() {
+    await Promise.resolve();
+}
+
+const g = async () => {
+    await f();
+}`,
+});