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

Author: rbuckton

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));
         }
@@ -16393,9 +16398,13 @@ 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;
         }
@@ -20796,6 +20805,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/contextuallyTypeAsyncFunctionReturnType.symbols

@@ -0,0 +1,29 @@
+=== 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))
+
+    });
+}

tests/baselines/reference/contextuallyTypeAsyncFunctionReturnType.types

@@ -0,0 +1,31 @@
+=== 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"
+
+    });
+}

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

@@ -0,0 +1,15 @@
+// @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" });
+    });
+}