Add contextual typing for await operand

Author: rbuckton

src/compiler/checker.ts

@@ -16409,6 +16409,15 @@ namespace ts {
             return undefined;
         }
 
+        function getContextualTypeForAwaitOperand(node: AwaitExpression): Type | undefined {
+            const contextualType = getContextualType(node);
+            if (contextualType) {
+                const contextualAwaitedType = getAwaitedType(contextualType);
+                return contextualAwaitedType && getUnionType([contextualAwaitedType, createPromiseLikeType(contextualAwaitedType)]);
+            }
+            return undefined;
+        }
+
         function getContextualTypeForYieldOperand(node: YieldExpression): Type | undefined {
             const func = getContainingFunction(node);
             if (func) {
@@ -16770,7 +16779,9 @@ namespace ts {
                     return getContextualTypeForReturnExpression(node);
                 case SyntaxKind.YieldExpression:
                     return getContextualTypeForYieldOperand(<YieldExpression>parent);
-                    case SyntaxKind.CallExpression:
+                case SyntaxKind.AwaitExpression:
+                    return getContextualTypeForAwaitOperand(<AwaitExpression>parent);
+                case SyntaxKind.CallExpression:
                 case SyntaxKind.NewExpression:
                     return getContextualTypeForArgument(<CallExpression | NewExpression>parent, node);
                 case SyntaxKind.TypeAssertionExpression:

tests/baselines/reference/contextuallyTypeAsyncFunctionReturnType.symbols

@@ -27,3 +27,28 @@ async function fn2(): Promise<Obj> {
 
     });
 }
+
+async function fn3(): Promise<Obj> {
+>fn3 : Symbol(fn3, Decl(contextuallyTypeAsyncFunctionReturnType.ts, 10, 1))
+>Promise : Symbol(Promise, Decl(lib.es5.d.ts, --, --), Decl(lib.es2015.iterable.d.ts, --, --), Decl(lib.es2015.promise.d.ts, --, --), Decl(lib.es2015.symbol.wellknown.d.ts, --, --), Decl(lib.es2018.promise.d.ts, --, --))
+>Obj : Symbol(Obj, Decl(contextuallyTypeAsyncFunctionReturnType.ts, 0, 0))
+
+    return await { key: "value" };
+>key : Symbol(key, Decl(contextuallyTypeAsyncFunctionReturnType.ts, 13, 18))
+}
+
+async function fn4(): Promise<Obj> {
+>fn4 : Symbol(fn4, Decl(contextuallyTypeAsyncFunctionReturnType.ts, 14, 1))
+>Promise : Symbol(Promise, Decl(lib.es5.d.ts, --, --), Decl(lib.es2015.iterable.d.ts, --, --), Decl(lib.es2015.promise.d.ts, --, --), Decl(lib.es2015.symbol.wellknown.d.ts, --, --), Decl(lib.es2018.promise.d.ts, --, --))
+>Obj : Symbol(Obj, Decl(contextuallyTypeAsyncFunctionReturnType.ts, 0, 0))
+
+    return await new Promise(resolve => {
+>Promise : Symbol(Promise, Decl(lib.es5.d.ts, --, --), Decl(lib.es2015.iterable.d.ts, --, --), Decl(lib.es2015.promise.d.ts, --, --), Decl(lib.es2015.symbol.wellknown.d.ts, --, --), Decl(lib.es2018.promise.d.ts, --, --))
+>resolve : Symbol(resolve, Decl(contextuallyTypeAsyncFunctionReturnType.ts, 17, 29))
+
+        resolve({ key: "value" });
+>resolve : Symbol(resolve, Decl(contextuallyTypeAsyncFunctionReturnType.ts, 17, 29))
+>key : Symbol(key, Decl(contextuallyTypeAsyncFunctionReturnType.ts, 18, 17))
+
+    });
+}

tests/baselines/reference/contextuallyTypeAsyncFunctionReturnType.types

@@ -29,3 +29,33 @@ async function fn2(): Promise<Obj> {
 
     });
 }
+
+async function fn3(): Promise<Obj> {
+>fn3 : () => Promise<Obj>
+
+    return await { key: "value" };
+>await { key: "value" } : { key: "value"; }
+>{ key: "value" } : { key: "value"; }
+>key : "value"
+>"value" : "value"
+}
+
+async function fn4(): Promise<Obj> {
+>fn4 : () => Promise<Obj>
+
+    return await new Promise(resolve => {
+>await new Promise(resolve => {        resolve({ key: "value" });    }) : Obj
+>new Promise(resolve => {        resolve({ key: "value" });    }) : Promise<Obj>
+>Promise : PromiseConstructor
+>resolve => {        resolve({ key: "value" });    } : (resolve: (value?: Obj | PromiseLike<Obj>) => void) => void
+>resolve : (value?: Obj | PromiseLike<Obj>) => void
+
+        resolve({ key: "value" });
+>resolve({ key: "value" }) : void
+>resolve : (value?: Obj | PromiseLike<Obj>) => void
+>{ key: "value" } : { key: "value"; }
+>key : "value"
+>"value" : "value"
+
+    });
+}

tests/cases/conformance/types/contextualTypes/asyncFunctions/contextuallyTypeAsyncFunctionReturnType.ts

@@ -12,4 +12,14 @@ async function fn2(): Promise<Obj> {
     return new Promise(resolve => {
         resolve({ key: "value" });
     });
+}
+
+async function fn3(): Promise<Obj> {
+    return await { key: "value" };
+}
+
+async function fn4(): Promise<Obj> {
+    return await new Promise(resolve => {
+        resolve({ key: "value" });
+    });
 }