feat: codefix for for await of (#50623)

microsoft · Sep 29, 2022 · 42f9143 · 42f9143
1 parent ecf50e8
commit 42f9143
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 1 deletion.
diff --git a/src/compiler/checker.ts b/src/compiler/checker.ts
@@ -642,6 +642,11 @@ namespace ts {
             getOptionalType: () => optionalType,
             getPromiseType: () => getGlobalPromiseType(/*reportErrors*/ false),
             getPromiseLikeType: () => getGlobalPromiseLikeType(/*reportErrors*/ false),
+            getAsyncIterableType: () => {
+                const type = getGlobalAsyncIterableType(/*reportErrors*/ false);
+                if (type === emptyGenericType) return undefined;
+                return type;
+            },
             isSymbolAccessible,
             isArrayType,
             isTupleType,
@@ -39429,7 +39434,18 @@ namespace ts {
             const message = allowAsyncIterables
                 ? Diagnostics.Type_0_must_have_a_Symbol_asyncIterator_method_that_returns_an_async_iterator
                 : Diagnostics.Type_0_must_have_a_Symbol_iterator_method_that_returns_an_iterator;
-            return errorAndMaybeSuggestAwait(errorNode, !!getAwaitedTypeOfPromise(type), message, typeToString(type));
+            const suggestAwait =
+                // for (const x of Promise<...>) or [...Promise<...>]
+                !!getAwaitedTypeOfPromise(type)
+                // for (const x of AsyncIterable<...>)
+                || (
+                    !allowAsyncIterables &&
+                    isForOfStatement(errorNode.parent) &&
+                    errorNode.parent.expression === errorNode &&
+                    getGlobalAsyncIterableType(/** reportErrors */ false) !== emptyGenericType &&
+                    isTypeAssignableTo(type, getGlobalAsyncIterableType(/** reportErrors */ false)
+                ));
+            return errorAndMaybeSuggestAwait(errorNode, suggestAwait, message, typeToString(type));
         }
 
         /**

diff --git a/src/compiler/types.ts b/src/compiler/types.ts
@@ -4760,6 +4760,7 @@ namespace ts {
         /* @internal */ createPromiseType(type: Type): Type;
         /* @internal */ getPromiseType(): Type;
         /* @internal */ getPromiseLikeType(): Type;
+        /* @internal */ getAsyncIterableType(): Type | undefined;
 
         /* @internal */ isTypeAssignableTo(source: Type, target: Type): boolean;
         /* @internal */ createAnonymousType(symbol: Symbol | undefined, members: SymbolTable, callSignatures: Signature[], constructSignatures: Signature[], indexInfos: IndexInfo[]): Type;

diff --git a/src/services/codefixes/addMissingAwait.ts b/src/services/codefixes/addMissingAwait.ts
@@ -226,6 +226,15 @@ namespace ts.codefix {
     }
 
     function makeChange(changeTracker: textChanges.ChangeTracker, errorCode: number, sourceFile: SourceFile, checker: TypeChecker, insertionSite: Expression, fixedDeclarations?: Set<number>) {
+        if (isForOfStatement(insertionSite.parent) && !insertionSite.parent.awaitModifier) {
+            const exprType = checker.getTypeAtLocation(insertionSite);
+            const asyncIter = checker.getAsyncIterableType();
+            if (asyncIter && checker.isTypeAssignableTo(exprType, asyncIter)) {
+                const forOf = insertionSite.parent;
+                changeTracker.replaceNode(sourceFile, forOf, factory.updateForOfStatement(forOf, factory.createToken(SyntaxKind.AwaitKeyword), forOf.initializer, forOf.expression, forOf.statement));
+                return;
+            }
+        }
         if (isBinaryExpression(insertionSite)) {
             for (const side of [insertionSite.left, insertionSite.right]) {
                 if (fixedDeclarations && isIdentifier(side)) {

diff --git a/tests/cases/fourslash/codeFixAddMissingAwait_forAwaitOf.ts b/tests/cases/fourslash/codeFixAddMissingAwait_forAwaitOf.ts
@@ -0,0 +1,17 @@
+/// <reference path="fourslash.ts" />
+// @lib: es2020
+// @target: es2020
+////async function* g() {}
+////async function fn() {
+////    for (const { } of g()) { }
+////}
+
+verify.codeFix({
+  description: ts.Diagnostics.Add_await.message,
+  index: 0,
+  newFileContent:
+`async function* g() {}
+async function fn() {
+    for await (const { } of g()) { }
+}`
+});