Support promise-like types in contextual return type of async function

src/compiler/checker.ts

@@ -495,6 +495,7 @@ namespace ts {
         let deferredGlobalESSymbolType: ObjectType;
         let deferredGlobalTypedPropertyDescriptorType: GenericType;
         let deferredGlobalPromiseType: GenericType;
+        let deferredGlobalPromiseLikeType: GenericType;
         let deferredGlobalPromiseConstructorSymbol: Symbol | undefined;
         let deferredGlobalPromiseConstructorLikeType: ObjectType;
         let deferredGlobalIterableType: GenericType;
@@ -8538,6 +8539,10 @@ namespace ts {
             return deferredGlobalPromiseType || (deferredGlobalPromiseType = getGlobalType("Promise" as __String, /*arity*/ 1, reportErrors)) || emptyGenericType;
         }
 
+        function getGlobalPromiseLikeType(reportErrors: boolean) {
+            return deferredGlobalPromiseLikeType || (deferredGlobalPromiseLikeType = getGlobalType("PromiseLike" as __String, /*arity*/ 1, reportErrors)) || emptyGenericType;
+        }
+
         function getGlobalPromiseConstructorSymbol(reportErrors: boolean): Symbol | undefined {
             return deferredGlobalPromiseConstructorSymbol || (deferredGlobalPromiseConstructorSymbol = getGlobalValueSymbol("Promise" as __String, reportErrors));
         }
@@ -16372,9 +16377,22 @@ namespace ts {
                 }
 
                 const contextualReturnType = getContextualReturnType(func);
-                return functionFlags & FunctionFlags.Async
-                    ? contextualReturnType && getAwaitedTypeOfPromise(contextualReturnType) // Async function
-                    : contextualReturnType; // Regular function
+                if (contextualReturnType) {
+                    if (functionFlags & FunctionFlags.Async) { // Async function
+                        const contextualAwaitedType = getAwaitedTypeOfPromise(contextualReturnType);
+                        return contextualAwaitedType && getUnionType([contextualAwaitedType, createPromiseLikeType(contextualAwaitedType)]);
+                    }
+                    return contextualReturnType; // Regular function
+                }
+            }
+            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;
         }
@@ -16740,7 +16758,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:
@@ -20775,6 +20795,18 @@ namespace ts {
             return emptyObjectType;
         }
 
+        function createPromiseLikeType(promisedType: Type): Type {
+            // creates a `PromiseLike<T>` type where `T` is the promisedType argument
+            const globalPromiseLikeType = getGlobalPromiseLikeType(/*reportErrors*/ true);
+            if (globalPromiseLikeType !== emptyGenericType) {
+                // if the promised type is itself a promise, get the underlying type; otherwise, fallback to the promised type
+                promisedType = getAwaitedType(promisedType) || emptyObjectType;
+                return createTypeReference(globalPromiseLikeType, [promisedType]);
+            }
+
+            return emptyObjectType;
+        }
+
         function createPromiseReturnType(func: FunctionLikeDeclaration | ImportCall, promisedType: Type) {
             const promiseType = createPromiseType(promisedType);
             if (promiseType === emptyObjectType) {

tests/baselines/reference/contextuallyTypeAsyncFunctionAwaitOperand.symbols

@@ -0,0 +1,27 @@
+=== tests/cases/conformance/types/contextualTypes/asyncFunctions/contextuallyTypeAsyncFunctionAwaitOperand.ts ===
+interface Obj { key: "value"; }
+>Obj : Symbol(Obj, Decl(contextuallyTypeAsyncFunctionAwaitOperand.ts, 0, 0))
+>key : Symbol(Obj.key, Decl(contextuallyTypeAsyncFunctionAwaitOperand.ts, 0, 15))
+
+async function fn1(): Promise<Obj> {
+>fn1 : Symbol(fn1, Decl(contextuallyTypeAsyncFunctionAwaitOperand.ts, 0, 31))
+>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(contextuallyTypeAsyncFunctionAwaitOperand.ts, 0, 0))
+
+    const obj1: Obj = await { key: "value" };
+>obj1 : Symbol(obj1, Decl(contextuallyTypeAsyncFunctionAwaitOperand.ts, 3, 9))
+>Obj : Symbol(Obj, Decl(contextuallyTypeAsyncFunctionAwaitOperand.ts, 0, 0))
+>key : Symbol(key, Decl(contextuallyTypeAsyncFunctionAwaitOperand.ts, 3, 29))
+
+    const obj2: Obj = await new Promise(resolve => resolve({ key: "value" }));
+>obj2 : Symbol(obj2, Decl(contextuallyTypeAsyncFunctionAwaitOperand.ts, 4, 9))
+>Obj : Symbol(Obj, Decl(contextuallyTypeAsyncFunctionAwaitOperand.ts, 0, 0))
+>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(contextuallyTypeAsyncFunctionAwaitOperand.ts, 4, 40))
+>resolve : Symbol(resolve, Decl(contextuallyTypeAsyncFunctionAwaitOperand.ts, 4, 40))
+>key : Symbol(key, Decl(contextuallyTypeAsyncFunctionAwaitOperand.ts, 4, 60))
+
+    return await { key: "value" };
+>key : Symbol(key, Decl(contextuallyTypeAsyncFunctionAwaitOperand.ts, 5, 18))
+}
+

tests/baselines/reference/contextuallyTypeAsyncFunctionAwaitOperand.types

@@ -0,0 +1,34 @@
+=== tests/cases/conformance/types/contextualTypes/asyncFunctions/contextuallyTypeAsyncFunctionAwaitOperand.ts ===
+interface Obj { key: "value"; }
+>key : "value"
+
+async function fn1(): Promise<Obj> {
+>fn1 : () => Promise<Obj>
+
+    const obj1: Obj = await { key: "value" };
+>obj1 : Obj
+>await { key: "value" } : { key: "value"; }
+>{ key: "value" } : { key: "value"; }
+>key : "value"
+>"value" : "value"
+
+    const obj2: Obj = await new Promise(resolve => resolve({ key: "value" }));
+>obj2 : Obj
+>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" }) : void
+>resolve : (value?: Obj | PromiseLike<Obj>) => void
+>{ key: "value" } : { key: "value"; }
+>key : "value"
+>"value" : "value"
+
+    return await { key: "value" };
+>await { key: "value" } : { key: "value"; }
+>{ key: "value" } : { key: "value"; }
+>key : "value"
+>"value" : "value"
+}
+

tests/baselines/reference/contextuallyTypeAsyncFunctionReturnType.symbols

@@ -0,0 +1,54 @@
+=== tests/cases/conformance/types/contextualTypes/asyncFunctions/contextuallyTypeAsyncFunctionReturnType.ts ===
+interface Obj { key: "value"; }
+>Obj : Symbol(Obj, Decl(contextuallyTypeAsyncFunctionReturnType.ts, 0, 0))
+>key : Symbol(Obj.key, Decl(contextuallyTypeAsyncFunctionReturnType.ts, 0, 15))
+
+async function fn1(): Promise<Obj> {
+>fn1 : Symbol(fn1, Decl(contextuallyTypeAsyncFunctionReturnType.ts, 0, 31))
+>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 { key: "value" };
+>key : Symbol(key, Decl(contextuallyTypeAsyncFunctionReturnType.ts, 3, 12))
+}
+
+async function fn2(): Promise<Obj> {
+>fn2 : Symbol(fn2, Decl(contextuallyTypeAsyncFunctionReturnType.ts, 4, 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 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, 7, 23))
+
+        resolve({ key: "value" });
+>resolve : Symbol(resolve, Decl(contextuallyTypeAsyncFunctionReturnType.ts, 7, 23))
+>key : Symbol(key, Decl(contextuallyTypeAsyncFunctionReturnType.ts, 8, 17))
+
+    });
+}
+
+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

@@ -0,0 +1,61 @@
+=== tests/cases/conformance/types/contextualTypes/asyncFunctions/contextuallyTypeAsyncFunctionReturnType.ts ===
+interface Obj { key: "value"; }
+>key : "value"
+
+async function fn1(): Promise<Obj> {
+>fn1 : () => Promise<Obj>
+
+    return { key: "value" };
+>{ key: "value" } : { key: "value"; }
+>key : "value"
+>"value" : "value"
+}
+
+async function fn2(): Promise<Obj> {
+>fn2 : () => Promise<Obj>
+
+    return new Promise(resolve => {
+>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"
+
+    });
+}
+
+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/contextuallyTypeAsyncFunctionAwaitOperand.ts

@@ -0,0 +1,11 @@
+// @target: esnext
+// @noImplicitAny: true
+// @noEmit: true
+
+interface Obj { key: "value"; }
+
+async function fn1(): Promise<Obj> {
+    const obj1: Obj = await { key: "value" };
+    const obj2: Obj = await new Promise(resolve => resolve({ key: "value" }));
+    return await { key: "value" };
+}

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

@@ -0,0 +1,25 @@
+// @target: esnext
+// @noImplicitAny: true
+// @noEmit: true
+
+interface Obj { key: "value"; }
+
+async function fn1(): Promise<Obj> {
+    return { key: "value" };
+}
+
+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" });
+    });
+}