Refactor conditional logic to avoid fallthrough and to make the contr…

…ol flow more clear
microsoft · Jan 31, 2022 · 0ea6851 · 0ea6851
1 parent 9d55f13
commit 0ea6851
Showing 1 changed file with 10 additions and 5 deletions.
diff --git a/src/compiler/checker.ts b/src/compiler/checker.ts
@@ -26033,7 +26033,7 @@ namespace ts {
         function getContextualTypeForReturnExpression(node: Expression): Type | undefined {
             const func = getContainingFunction(node);
             if (func) {
-                let contextualReturnType = getContextualReturnType(func);
+                const contextualReturnType = getContextualReturnType(func);
                 if (contextualReturnType) {
                     const functionFlags = getFunctionFlags(func);
                     if (functionFlags & FunctionFlags.Generator) { // Generator or AsyncGenerator function
@@ -26042,12 +26042,17 @@ namespace ts {
                         if (!iterationTypes) {
                             return undefined;
                         }
-                        contextualReturnType = iterationTypes.returnType;
+                        if (functionFlags & FunctionFlags.Async) {
+                            // unwrap Promise to get the awaited type without the `Awaited<T>` alias
+                            const contextualAwaitedType = mapType(iterationTypes.returnType, getAwaitedTypeNoAlias);
+                            return contextualAwaitedType && getUnionType([contextualAwaitedType, createPromiseLikeType(contextualAwaitedType)]);
+                        }
+                        return iterationTypes.returnType;
                     }
 
-                    if (functionFlags & FunctionFlags.Async) { // Async function or AsyncGenerator function
-                        const contextualAwaitedType = mapType(contextualReturnType, functionFlags & FunctionFlags.Generator ? getAwaitedTypeNoAlias : getAwaitedTypeOfPromise);
-                        return contextualAwaitedType && getUnionType([contextualAwaitedType, createPromiseLikeType(contextualAwaitedType)]);
+                    if (functionFlags & FunctionFlags.Async) {
+                        const contextualTypeOfPromise = mapType(contextualReturnType, getAwaitedTypeOfPromise);
+                        return contextualTypeOfPromise && getUnionType([contextualTypeOfPromise, createPromiseLikeType(contextualTypeOfPromise)]);
                     }
 
                     return contextualReturnType; // Regular function or Generator function