[WIP] Promises to Async Await Code Fix

src/compiler/checker.ts

@@ -289,6 +289,7 @@ namespace ts {
             getNeverType: () => neverType,
             isSymbolAccessible,
             isArrayLikeType,
+            isPromiseLikeType,
             getAllPossiblePropertiesOfTypes,
             getSuggestionForNonexistentProperty: (node, type) => getSuggestionForNonexistentProperty(node, type),
             getSuggestionForNonexistentSymbol: (location, name, meaning) => getSuggestionForNonexistentSymbol(location, escapeLeadingUnderscores(name), meaning),
@@ -419,6 +420,7 @@ namespace ts {
         let globalObjectType: ObjectType;
         let globalFunctionType: ObjectType;
         let globalArrayType: GenericType;
+        let globalPromiseType: GenericType;
         let globalReadonlyArrayType: GenericType;
         let globalStringType: ObjectType;
         let globalNumberType: ObjectType;
@@ -11810,6 +11812,11 @@ namespace ts {
             return getObjectFlags(type) & ObjectFlags.Reference && (<TypeReference>type).target === globalArrayType;
         }
 
+        function isPromiseLikeType(type: Type): boolean {
+            return getObjectFlags(type) & ObjectFlags.Reference && (<TypeReference>type).target === globalPromiseType /*|| (<TypeReference>type).target === globalReadonlyArrayType)*/ ||
+                !(type.flags & TypeFlags.Nullable) && isTypeAssignableTo(type, globalPromiseType);
+        }
+
         function isArrayLikeType(type: Type): boolean {
             // A type is array-like if it is a reference to the global Array or global ReadonlyArray type,
             // or if it is not the undefined or null type and if it is assignable to ReadonlyArray<any>
@@ -27028,6 +27035,7 @@ namespace ts {
             getSymbolLinks(unknownSymbol).type = unknownType;
 
             // Initialize special types
+            globalPromiseType = getGlobalPromiseType(/*reportErrors*/ true);
             globalArrayType = getGlobalType("Array" as __String, /*arity*/ 1, /*reportErrors*/ true);
             globalObjectType = getGlobalType("Object" as __String, /*arity*/ 0, /*reportErrors*/ true);
             globalFunctionType = getGlobalType("Function" as __String, /*arity*/ 0, /*reportErrors*/ true);

src/compiler/diagnosticMessages.json

@@ -3964,6 +3964,10 @@
         "category": "Suggestion",
         "code": 80005
     },
+    "This may be converted to use async and await.": {
+        "category": "Suggestion",
+        "code": 80006
+    },
 
     "Add missing 'super()' call": {
         "category": "Message",
@@ -4280,5 +4284,13 @@
     "Remove all unused labels": {
         "category": "Message",
         "code": 95054
+    },
+    "Convert to use async and await":{
+        "category": "Message",
+         "code": 95055
+    },
+    "Convert all to use async and await": {
+        "category": "Message",
+        "code": 95056
     }
 }

src/compiler/types.ts

@@ -3046,6 +3046,7 @@ namespace ts {
          * Does not include properties of primitive types.
          */
         /* @internal */ isArrayLikeType(type: Type): boolean;
+        /* @internal */ isPromiseLikeType(type: Type): boolean;
         /* @internal */ getAllPossiblePropertiesOfTypes(type: ReadonlyArray<Type>): Symbol[];
         /* @internal */ resolveName(name: string, location: Node, meaning: SymbolFlags, excludeGlobals: boolean): Symbol | undefined;
         /* @internal */ getJsxNamespace(location?: Node): string;

src/harness/tsconfig.json

@@ -93,6 +93,7 @@
         "../services/codefixes/annotateWithTypeFromJSDoc.ts",
         "../services/codefixes/convertFunctionToEs6Class.ts",
         "../services/codefixes/convertToEs6Module.ts",
+        "../services/codefixes/convertPromisesToAwaitAndAsync.ts",
         "../services/codefixes/correctQualifiedNameToIndexedAccessType.ts",
         "../services/codefixes/fixClassIncorrectlyImplementsInterface.ts",
         "../services/codefixes/importFixes.ts",

src/services/codefixes/convertPromisesToAwaitAndAsync.ts

@@ -0,0 +1,128 @@
+namespace ts.codefix {
+    const fixId = "convertPromisesToAwaitAndAsync";
+    const errorCodes = [Diagnostics.This_may_be_converted_to_use_async_and_await.code];
+    registerCodeFix({
+        errorCodes,
+        getCodeActions(context: CodeFixContext) {
+            const changes = textChanges.ChangeTracker.with(context, (t) => convertToAsyncAwait(t, context.sourceFile, context.span.start, context.program.getTypeChecker()));
+            return [createCodeFixAction(fixId, changes, Diagnostics.Convert_to_use_async_and_await, fixId, Diagnostics.Convert_all_to_use_async_and_await)];
+        },
+        fixIds: [fixId],
+        getAllCodeActions: context => codeFixAll(context, errorCodes, (changes, err) => convertToAsyncAwait(changes, err.file!, err.start, context.program.getTypeChecker())),
+    });
+    function convertToAsyncAwait(changes: textChanges.ChangeTracker, sourceFile: SourceFile, position: number, checker: TypeChecker): void {
+        // get the function declaration - returns a promise
+        const funcToConvert: FunctionLikeDeclaration = checker.getSymbolAtLocation(getTokenAtPosition(sourceFile, position, /*includeEndPosition*/ false)).valueDeclaration as FunctionLikeDeclaration;
+        // add the async keyword
+        changes.insertModifierBefore(sourceFile, SyntaxKind.AsyncKeyword, funcToConvert);
+
+        const stmts: NodeArray<Statement> = (<FunctionBody>funcToConvert.body).statements;
+        for (const stmt of stmts) {
+            forEachChild(stmt, function visit(node: Node) {
+                if (node.kind === SyntaxKind.CallExpression) {
+                    const newNode = parseCallback(node as CallExpression, checker);
+                    if (newNode) {
+                        changes.replaceNodeWithNodes(sourceFile, stmt, newNode);
+                    }
+                }
+                else if (!isFunctionLike(node)) {
+                    forEachChild(node, visit);
+                }
+            });
+        }
+    }
+
+    function parseCallback(node: Expression, checker: TypeChecker, argName?: string, argUsed?: boolean): Statement[] {
+        if (!node) {
+            return;
+        }
+
+        if (node.kind === SyntaxKind.CallExpression && checker.isPromiseLikeType(checker.getTypeAtLocation(node)) && (<CallExpression>node).expression.kind !== SyntaxKind.PropertyAccessExpression) {
+            return parsePromiseCall(node as CallExpression, argName, argUsed);
+        }
+        else if (node.kind === SyntaxKind.CallExpression && isCallback(node as CallExpression, "then", checker)) {
+            return parseThen(node as CallExpression, checker);
+        }
+        else if (node.kind === SyntaxKind.CallExpression && isCallback(node as CallExpression, "catch", checker)) {
+            return parseCatch(node as CallExpression, checker);
+        }
+        else if (node.kind === SyntaxKind.PropertyAccessExpression) {
+            return parseCallback((<PropertyAccessExpression>node).expression, checker, argName, argUsed);
+        }
+    }
+
+    function parseCatch(node: CallExpression, checker: TypeChecker): Statement[] {
+        const func = node.arguments[0];
+        const argName = getArgName(func, "arg", checker);
+        const tryBlock = createBlock(parseCallback(node.expression, checker, argName, argName !== "arg"));
+        // instead of using e -> get the paramater of the catch function and use that
+        const catchClause = createCatchClause(argName, createBlock(getCallbackBody(func, argName)));
+        return [createTry(tryBlock, catchClause, /*finallyBlock*/ undefined)];
+    }
+
+    function parseThen(node: CallExpression, checker: TypeChecker): Statement[] {
+        const res = node.arguments[0];
+        const rej = node.arguments[1];
+        // TODO - what if this is a binding pattern and not an Identifier
+        const argNameRes = getArgName(res, "val", checker);
+
+        if (rej) {
+            const argNameRej = getArgName(rej, "e", checker);
+
+            const tryBlock = createBlock(parseCallback(node.expression, checker, argNameRes, argNameRes !== "val"));
+            const catchClause = createCatchClause(argNameRej, createBlock(getCallbackBody(rej, argNameRej)));
+
+            return [createTry(tryBlock, catchClause, /*finalllyBlock*/ undefined) as Statement].concat(getCallbackBody(res, argNameRes));
+        }
+        else {
+            return parseCallback(node.expression, checker, argNameRes, argNameRes !== "val").concat(getCallbackBody(res, argNameRes));
+        }
+    }
+
+    function parsePromiseCall(node: CallExpression, argName?: string, argUsed?: boolean): Statement[] {
+        if (!argName) {
+            argName = "val"; // fix this to maybe not always create a variable declaration if not necessary
+        }
+        return argUsed ? [createVariableStatement(/*modifiers*/ undefined, [createVariableDeclaration(createIdentifier(argName), /*type*/ undefined, createAwait(node))])] :
+                         [createStatement(createAwait(node))];
+    }
+
+    function getCallbackBody(func: Node, argName: string): NodeArray<Statement> {
+        switch (func.kind) {
+            case SyntaxKind.Identifier:
+                return createNodeArray([(createReturn(createAwait(createCall(func as Identifier, /*typeArguments*/ undefined, [createIdentifier(argName)]))))]);
+            case SyntaxKind.FunctionDeclaration:
+            case SyntaxKind.ArrowFunction:
+                if ((<FunctionDeclaration>func).body.kind === SyntaxKind.Block) {
+                    return (<FunctionDeclaration>func).body.statements;
+                }
+                else {
+                    return createNodeArray([createStatement((<ArrowFunction>func).body as Expression)]);
+                }
+        }
+    }
+    function isCallback(node: CallExpression, funcName: string, checker: TypeChecker): boolean {
+        if (node.expression.kind !== SyntaxKind.PropertyAccessExpression) {
+            return false;
+        }
+        return (<PropertyAccessExpression>node.expression).name.text === funcName && checker.isPromiseLikeType(checker.getTypeAtLocation(node));
+    }
+    function getArgName(funcNode: Node, defaultVal: string, checker: TypeChecker): string {
+        if (isFunctionLikeDeclaration(funcNode) && funcNode.parameters.length > 0) {
+            const name = (<Identifier>funcNode.parameters[0].name).text;
+            if (name !== "_") {
+                return name;
+            } 
+            else {
+                return defaultVal;
+            }
+        }
+        else if (checker.getTypeAtLocation(funcNode).getCallSignatures().length > 0 && checker.getTypeAtLocation(funcNode).getCallSignatures()[0].parameters.length > 0) {
+            const name = checker.getTypeAtLocation(funcNode).getCallSignatures()[0].parameters[0].name;
+            return name;
+        }
+        else {
+            return defaultVal;
+        }
+    }
+}

src/services/suggestionDiagnostics.ts

@@ -17,12 +17,30 @@ namespace ts {
             switch (node.kind) {
                 case SyntaxKind.FunctionDeclaration:
                 case SyntaxKind.FunctionExpression:
+
                     if (isJsFile) {
                         const symbol = node.symbol;
                         if (symbol.members && (symbol.members.size > 0)) {
                             diags.push(createDiagnosticForNode(isVariableDeclaration(node.parent) ? node.parent.name : node, Diagnostics.This_constructor_function_may_be_converted_to_a_class_declaration));
                         }
                     }
+                    if (node.modifiers && isAsyncFunction(node)) {
+                        break;
+                    }
+
+                    const returnType = checker.getReturnTypeOfSignature(checker.getSignatureFromDeclaration(<FunctionDeclaration | FunctionExpression>node));
+                    if (!returnType || !returnType.symbol) {
+                        break;
+                    }
+
+                    // collect all the return statements
+                    // check that a property access expression exists in there and that it is a handler
+                    const retStmts: ReturnStatement[] = [];
+                    getReturnStmts(node, retStmts);
+
+                    if (isPromiseType(returnType) && retStmts.length > 0) {
+                        diags.push(createDiagnosticForNode(isVariableDeclaration(node.parent) ? node.parent.name : node, Diagnostics.This_may_be_converted_to_use_async_and_await));
+                    }
                     break;
             }
 
@@ -104,4 +122,34 @@ namespace ts {
     function getErrorNodeFromCommonJsIndicator(commonJsModuleIndicator: Node): Node {
         return isBinaryExpression(commonJsModuleIndicator) ? commonJsModuleIndicator.left : commonJsModuleIndicator;
     }
+
+    function isPromiseType(T: Type): boolean {
+        return T.flags === TypeFlags.Object && T.symbol.name === "Promise";
+    }
+
+    function getReturnStmts(node: Node, retStmts: ReturnStatement[]) {
+        forEachChild(node, visit);
+
+        function visit(node: Node) {
+            if (isFunctionLike(node)) {
+                return;
+            }
+
+            if (node.kind === SyntaxKind.ReturnStatement) {
+                forEachChild(node, hasCallback);
+            }
+
+            function hasCallback(child: Node) {
+                if (child.kind === SyntaxKind.CallExpression && (<CallExpression>child).expression.kind === SyntaxKind.PropertyAccessExpression &&
+                    ((<PropertyAccessExpression>(<CallExpression>child).expression).name.text === "then" || (<PropertyAccessExpression>(<CallExpression>child).expression).name.text === "catch")) {
+                    retStmts.push(node as ReturnStatement);
+                }
+                else if (!isFunctionLike(child)) {
+                    forEachChild(child, hasCallback);
+                }
+            }
+
+            forEachChild(node, visit);
+        }
+    }
 }

tests/cases/fourslash/codeFixPromiseToAsync.ts

@@ -0,0 +1,22 @@
+/// <reference path='fourslash.ts' />
+
+// @target: es6
+
+////function [|f|](): Promise<void> {
+////    return fetch('http://yahoo.com').then(result => { console.log(result); });
+////}
+
+verify.getSuggestionDiagnostics([{
+    message: "This may be converted to use async and await.",
+    code: 80006,
+}]);
+
+verify.codeFix({
+    description: "Convert to use async and await",
+    index: 0,
+    newFileContent:
+`async function f(): Promise<void> {
+    var result = await fetch('http://yahoo.com');
+    console.log(result);
+}`,
+});

tests/cases/fourslash/codeFixPromiseToAsyncCatch.ts

@@ -0,0 +1,28 @@
+/// <reference path='fourslash.ts' />
+
+// @target: es6
+
+////function [|f|]():Promise<void> {
+////    return fetch('http://yahoo.com').then(result => { console.log(result); }).catch(err => { console.log(err); });
+////}
+
+verify.getSuggestionDiagnostics([{
+    message: "This may be converted to use async and await.",
+    code: 80006,
+}]);
+
+
+verify.codeFix({
+    description: "Convert to use async and await",
+    index: 0,
+    newFileContent:
+`async function f():Promise<void> {
+    try {
+        var result = await fetch('http://yahoo.com');
+        console.log(result);
+    }
+    catch (err) {
+        console.log(err);
+    }
+}`,
+});

tests/cases/fourslash/codeFixPromiseToAsyncCatchAndRej.ts

@@ -0,0 +1,33 @@
+/// <reference path='fourslash.ts' />
+
+// @target: es6
+
+////function [|f|]():Promise<void> {
+////    return fetch('http://yahoo.com').then(result => { console.log(result); }, rejection => { console.log("rejected:", rejection); }).catch(err => { console.log(err) });
+////}
+
+verify.getSuggestionDiagnostics([{
+    message: "This may be converted to use async and await.",
+    code: 80006,
+}]);
+
+
+verify.codeFix({
+    description: "Convert to use async and await",
+    index: 0,
+    newFileContent:
+`async function f():Promise<void> {
+    try {
+        try {
+            var result = await fetch('http://yahoo.com');
+        }
+        catch (rejection) {
+            console.log("rejected:", rejection);
+        }
+        console.log(result);
+    }
+    catch (err) {
+        console.log(err);
+    }
+}`,
+});