Add quick fix to add missing 'await' (#32356)

* Start prototyping addMissingAwait codefix

* Filter by diagnostics that have missing-await related info

* Start writing tests and checking precedence

* Implement codeFixAll, add test for binary expressions

* Add test for iterables

* Add test for passing argument

* Add test for call/construct signatures

* Add test for awaiting initializer

* Improve assertion error

* Replace specific property access error with general one and add await related info

* Add test for property access

* Require code to be inside a function body to offer await

* Accept suggestion

Co-Authored-By: Nathan Shively-Sanders <293473+sandersn@users.noreply.github.com>

* Add explicit test for code fix being not available unless something is a Promise

* Skip looking for function body if already in AwaitContext flags

* Inline getCodeActions function for symmetry

src/compiler/checker.ts

@@ -20715,7 +20715,8 @@ namespace ts {
             else {
                 const promisedType = getPromisedTypeOfPromise(containingType);
                 if (promisedType && getPropertyOfType(promisedType, propNode.escapedText)) {
-                    errorInfo = chainDiagnosticMessages(errorInfo, Diagnostics.Property_0_does_not_exist_on_type_1_Did_you_forget_to_use_await, declarationNameToString(propNode), typeToString(containingType));
+                    errorInfo = chainDiagnosticMessages(errorInfo, Diagnostics.Property_0_does_not_exist_on_type_1, declarationNameToString(propNode), typeToString(containingType));
+                    relatedInfo = createDiagnosticForNode(propNode, Diagnostics.Did_you_forget_to_use_await);
                 }
                 else {
                     const suggestion = getSuggestedSymbolForNonexistentProperty(propNode, containingType);

src/compiler/diagnosticMessages.json

@@ -2076,10 +2076,6 @@
         "category": "Error",
         "code": 2569
     },
-    "Property '{0}' does not exist on type '{1}'. Did you forget to use 'await'?": {
-        "category": "Error",
-        "code": 2570
-    },
     "Object is of type 'unknown'.": {
         "category": "Error",
         "code": 2571
@@ -5087,6 +5083,19 @@
         "category": "Message",
         "code": 95082
     },
+    "Add 'await'": {
+        "category": "Message",
+        "code": 95083
+    },
+    "Add 'await' to initializer for '{0}'": {
+        "category": "Message",
+        "code": 95084
+    },
+    "Fix all expressions possibly missing 'await'": {
+        "category": "Message",
+        "code": 95085
+    },
+
     "No value exists in scope for the shorthand property '{0}'. Either declare one or provide an initializer.": {
         "category": "Error",
         "code": 18004

src/harness/fourslash.ts

@@ -2828,11 +2828,28 @@ Actual: ${stringify(fullActual)}`);
             }
         }
 
-        public verifyCodeFixAvailable(negative: boolean, expected: FourSlashInterface.VerifyCodeFixAvailableOptions[] | undefined): void {
-            assert(!negative || !expected);
+        public verifyCodeFixAvailable(negative: boolean, expected: FourSlashInterface.VerifyCodeFixAvailableOptions[] | string | undefined): void {
             const codeFixes = this.getCodeFixes(this.activeFile.fileName);
-            const actuals = codeFixes.map((fix): FourSlashInterface.VerifyCodeFixAvailableOptions => ({ description: fix.description, commands: fix.commands }));
-            this.assertObjectsEqual(actuals, negative ? ts.emptyArray : expected);
+            if (negative) {
+                if (typeof expected === "undefined") {
+                    this.assertObjectsEqual(codeFixes, ts.emptyArray);
+                }
+                else if (typeof expected === "string") {
+                    if (codeFixes.some(fix => fix.fixName === expected)) {
+                        this.raiseError(`Expected not to find a fix with the name '${expected}', but one exists.`);
+                    }
+                }
+                else {
+                    assert(typeof expected === "undefined" || typeof expected === "string", "With a negated assertion, 'expected' must be undefined or a string value of a codefix name.");
+                }
+            }
+            else if (typeof expected === "string") {
+                this.assertObjectsEqual(codeFixes.map(fix => fix.fixName), [expected]);
+            }
+            else {
+                const actuals = codeFixes.map((fix): FourSlashInterface.VerifyCodeFixAvailableOptions => ({ description: fix.description, commands: fix.commands }));
+                this.assertObjectsEqual(actuals, negative ? ts.emptyArray : expected);
+            }
         }
 
         public verifyApplicableRefactorAvailableAtMarker(negative: boolean, markerName: string) {

src/services/codefixes/addMissingAwait.ts

@@ -0,0 +1,173 @@
+/* @internal */
+namespace ts.codefix {
+    type ContextualTrackChangesFunction = (cb: (changeTracker: textChanges.ChangeTracker) => void) => FileTextChanges[];
+    const fixId = "addMissingAwait";
+    const propertyAccessCode = Diagnostics.Property_0_does_not_exist_on_type_1.code;
+    const callableConstructableErrorCodes = [
+        Diagnostics.This_expression_is_not_callable.code,
+        Diagnostics.This_expression_is_not_constructable.code,
+    ];
+    const errorCodes = [
+        Diagnostics.An_arithmetic_operand_must_be_of_type_any_number_bigint_or_an_enum_type.code,
+        Diagnostics.The_left_hand_side_of_an_arithmetic_operation_must_be_of_type_any_number_bigint_or_an_enum_type.code,
+        Diagnostics.The_right_hand_side_of_an_arithmetic_operation_must_be_of_type_any_number_bigint_or_an_enum_type.code,
+        Diagnostics.Operator_0_cannot_be_applied_to_type_1.code,
+        Diagnostics.Operator_0_cannot_be_applied_to_types_1_and_2.code,
+        Diagnostics.This_condition_will_always_return_0_since_the_types_1_and_2_have_no_overlap.code,
+        Diagnostics.Type_0_is_not_an_array_type.code,
+        Diagnostics.Type_0_is_not_an_array_type_or_a_string_type.code,
+        Diagnostics.Type_0_is_not_an_array_type_or_a_string_type_Use_compiler_option_downlevelIteration_to_allow_iterating_of_iterators.code,
+        Diagnostics.Type_0_is_not_an_array_type_or_a_string_type_or_does_not_have_a_Symbol_iterator_method_that_returns_an_iterator.code,
+        Diagnostics.Type_0_is_not_an_array_type_or_does_not_have_a_Symbol_iterator_method_that_returns_an_iterator.code,
+        Diagnostics.Type_0_must_have_a_Symbol_iterator_method_that_returns_an_iterator.code,
+        Diagnostics.Type_0_must_have_a_Symbol_asyncIterator_method_that_returns_an_async_iterator.code,
+        Diagnostics.Argument_of_type_0_is_not_assignable_to_parameter_of_type_1.code,
+        propertyAccessCode,
+        ...callableConstructableErrorCodes,
+    ];
+
+    registerCodeFix({
+        fixIds: [fixId],
+        errorCodes,
+        getCodeActions: context => {
+            const { sourceFile, errorCode, span, cancellationToken, program } = context;
+            const expression = getAwaitableExpression(sourceFile, errorCode, span, cancellationToken, program);
+            if (!expression) {
+                return;
+            }
+
+            const checker = context.program.getTypeChecker();
+            const trackChanges: ContextualTrackChangesFunction = cb => textChanges.ChangeTracker.with(context, cb);
+            return compact([
+                getDeclarationSiteFix(context, expression, errorCode, checker, trackChanges),
+                getUseSiteFix(context, expression, errorCode, checker, trackChanges)]);
+        },
+        getAllCodeActions: context => {
+            const { sourceFile, program, cancellationToken } = context;
+            const checker = context.program.getTypeChecker();
+            return codeFixAll(context, errorCodes, (t, diagnostic) => {
+                const expression = getAwaitableExpression(sourceFile, diagnostic.code, diagnostic, cancellationToken, program);
+                if (!expression) {
+                    return;
+                }
+                const trackChanges: ContextualTrackChangesFunction = cb => (cb(t), []);
+                return getDeclarationSiteFix(context, expression, diagnostic.code, checker, trackChanges)
+                    || getUseSiteFix(context, expression, diagnostic.code, checker, trackChanges);
+            });
+        },
+    });
+
+    function getDeclarationSiteFix(context: CodeFixContext | CodeFixAllContext, expression: Expression, errorCode: number, checker: TypeChecker, trackChanges: ContextualTrackChangesFunction) {
+        const { sourceFile } = context;
+        const awaitableInitializer = findAwaitableInitializer(expression, sourceFile, checker);
+        if (awaitableInitializer) {
+            const initializerChanges = trackChanges(t => makeChange(t, errorCode, sourceFile, checker, awaitableInitializer));
+            return createCodeFixActionNoFixId(
+                "addMissingAwaitToInitializer",
+                initializerChanges,
+                [Diagnostics.Add_await_to_initializer_for_0, expression.getText(sourceFile)]);
+        }
+    }
+
+    function getUseSiteFix(context: CodeFixContext | CodeFixAllContext, expression: Expression, errorCode: number, checker: TypeChecker, trackChanges: ContextualTrackChangesFunction) {
+        const changes = trackChanges(t => makeChange(t, errorCode, context.sourceFile, checker, expression));
+        return createCodeFixAction(fixId, changes, Diagnostics.Add_await, fixId, Diagnostics.Fix_all_expressions_possibly_missing_await);
+    }
+
+    function isMissingAwaitError(sourceFile: SourceFile, errorCode: number, span: TextSpan, cancellationToken: CancellationToken, program: Program) {
+        const checker = program.getDiagnosticsProducingTypeChecker();
+        const diagnostics = checker.getDiagnostics(sourceFile, cancellationToken);
+        return some(diagnostics, ({ start, length, relatedInformation, code }) =>
+            isNumber(start) && isNumber(length) && textSpansEqual({ start, length }, span) &&
+            code === errorCode &&
+            !!relatedInformation &&
+            some(relatedInformation, related => related.code === Diagnostics.Did_you_forget_to_use_await.code));
+    }
+
+    function getAwaitableExpression(sourceFile: SourceFile, errorCode: number, span: TextSpan, cancellationToken: CancellationToken, program: Program): Expression | undefined {
+        const token = getTokenAtPosition(sourceFile, span.start);
+        // Checker has already done work to determine that await might be possible, and has attached
+        // related info to the node, so start by finding the expression that exactly matches up
+        // with the diagnostic range.
+        const expression = findAncestor(token, node => {
+            if (node.getStart(sourceFile) < span.start || node.getEnd() > textSpanEnd(span)) {
+                return "quit";
+            }
+            return isExpression(node) && textSpansEqual(span, createTextSpanFromNode(node, sourceFile));
+        }) as Expression | undefined;
+
+        return expression
+            && isMissingAwaitError(sourceFile, errorCode, span, cancellationToken, program)
+            && isInsideAwaitableBody(expression)
+                ? expression
+                : undefined;
+    }
+
+    function findAwaitableInitializer(expression: Node, sourceFile: SourceFile, checker: TypeChecker): Expression | undefined {
+        if (!isIdentifier(expression)) {
+            return;
+        }
+
+        const symbol = checker.getSymbolAtLocation(expression);
+        if (!symbol) {
+            return;
+        }
+
+        const declaration = tryCast(symbol.valueDeclaration, isVariableDeclaration);
+        const variableName = tryCast(declaration && declaration.name, isIdentifier);
+        const variableStatement = getAncestor(declaration, SyntaxKind.VariableStatement);
+        if (!declaration || !variableStatement ||
+            declaration.type ||
+            !declaration.initializer ||
+            variableStatement.getSourceFile() !== sourceFile ||
+            hasModifier(variableStatement, ModifierFlags.Export) ||
+            !variableName ||
+            !isInsideAwaitableBody(declaration.initializer)) {
+            return;
+        }
+
+        const isUsedElsewhere = FindAllReferences.Core.eachSymbolReferenceInFile(variableName, checker, sourceFile, identifier => {
+            return identifier !== expression;
+        });
+
+        if (isUsedElsewhere) {
+            return;
+        }
+
+        return declaration.initializer;
+    }
+
+    function isInsideAwaitableBody(node: Node) {
+        return node.kind & NodeFlags.AwaitContext || !!findAncestor(node, ancestor =>
+            ancestor.parent && isArrowFunction(ancestor.parent) && ancestor.parent.body === ancestor ||
+            isBlock(ancestor) && (
+                ancestor.parent.kind === SyntaxKind.FunctionDeclaration ||
+                ancestor.parent.kind === SyntaxKind.FunctionExpression ||
+                ancestor.parent.kind === SyntaxKind.ArrowFunction ||
+                ancestor.parent.kind === SyntaxKind.MethodDeclaration));
+    }
+
+    function makeChange(changeTracker: textChanges.ChangeTracker, errorCode: number, sourceFile: SourceFile, checker: TypeChecker, insertionSite: Expression) {
+        if (isBinaryExpression(insertionSite)) {
+            const { left, right } = insertionSite;
+            const leftType = checker.getTypeAtLocation(left);
+            const rightType = checker.getTypeAtLocation(right);
+            const newLeft = checker.getPromisedTypeOfPromise(leftType) ? createAwait(left) : left;
+            const newRight = checker.getPromisedTypeOfPromise(rightType) ? createAwait(right) : right;
+            changeTracker.replaceNode(sourceFile, left, newLeft);
+            changeTracker.replaceNode(sourceFile, right, newRight);
+        }
+        else if (errorCode === propertyAccessCode && isPropertyAccessExpression(insertionSite.parent)) {
+            changeTracker.replaceNode(
+                sourceFile,
+                insertionSite.parent.expression,
+                createParen(createAwait(insertionSite.parent.expression)));
+        }
+        else if (contains(callableConstructableErrorCodes, errorCode) && isCallOrNewExpression(insertionSite.parent)) {
+            changeTracker.replaceNode(sourceFile, insertionSite, createParen(createAwait(insertionSite)));
+        }
+        else {
+            changeTracker.replaceNode(sourceFile, insertionSite, createAwait(insertionSite));
+        }
+    }
+}

src/services/textChanges.ts

@@ -704,6 +704,10 @@ namespace ts.textChanges {
             }
         }
 
+        public parenthesizeExpression(sourceFile: SourceFile, expression: Expression) {
+            this.replaceRange(sourceFile, rangeOfNode(expression), createParen(expression));
+        }
+
         private finishClassesWithNodesInsertedAtStart(): void {
             this.classesWithNodesInsertedAtStart.forEach(({ node, sourceFile }) => {
                 const [openBraceEnd, closeBraceEnd] = getClassOrObjectBraceEnds(node, sourceFile);

src/services/tsconfig.json

@@ -45,6 +45,7 @@
         "codeFixProvider.ts",
         "refactorProvider.ts",
         "codefixes/addConvertToUnknownForNonOverlappingTypes.ts",
+        "codefixes/addMissingAwait.ts",
         "codefixes/addMissingConst.ts",
         "codefixes/addMissingInvocationForDecorator.ts",
         "codefixes/addNameToNamelessParameter.ts",

tests/baselines/reference/operationsAvailableOnPromisedType.errors.txt

@@ -8,7 +8,7 @@ tests/cases/compiler/operationsAvailableOnPromisedType.ts(17,5): error TS2367: T
 tests/cases/compiler/operationsAvailableOnPromisedType.ts(18,9): error TS2461: Type 'Promise<string[]>' is not an array type.
 tests/cases/compiler/operationsAvailableOnPromisedType.ts(19,21): error TS2495: Type 'Promise<string[]>' is not an array type or a string type.
 tests/cases/compiler/operationsAvailableOnPromisedType.ts(20,12): error TS2345: Argument of type 'Promise<number>' is not assignable to parameter of type 'number'.
-tests/cases/compiler/operationsAvailableOnPromisedType.ts(21,11): error TS2570: Property 'prop' does not exist on type 'Promise<{ prop: string; }>'. Did you forget to use 'await'?
+tests/cases/compiler/operationsAvailableOnPromisedType.ts(21,11): error TS2339: Property 'prop' does not exist on type 'Promise<{ prop: string; }>'.
 tests/cases/compiler/operationsAvailableOnPromisedType.ts(23,27): error TS2495: Type 'Promise<string[]>' is not an array type or a string type.
 tests/cases/compiler/operationsAvailableOnPromisedType.ts(24,5): error TS2349: This expression is not callable.
   Type 'Promise<() => void>' has no call signatures.
@@ -72,7 +72,8 @@ tests/cases/compiler/operationsAvailableOnPromisedType.ts(27,5): error TS2349: T
 !!! related TS2773 tests/cases/compiler/operationsAvailableOnPromisedType.ts:20:12: Did you forget to use 'await'?
             d.prop;
               ~~~~
-!!! error TS2570: Property 'prop' does not exist on type 'Promise<{ prop: string; }>'. Did you forget to use 'await'?
+!!! error TS2339: Property 'prop' does not exist on type 'Promise<{ prop: string; }>'.
+!!! related TS2773 tests/cases/compiler/operationsAvailableOnPromisedType.ts:21:11: Did you forget to use 'await'?
         }
         for await (const s of c) {}
                               ~

tests/cases/fourslash/codeFixAddMissingAwait_argument.ts

@@ -0,0 +1,13 @@
+/// <reference path="fourslash.ts" />
+////async function fn(a: Promise<string>, b: string) {
+////  fn(a, a);
+////}
+
+verify.codeFix({
+  description: ts.Diagnostics.Add_await.message,
+  index: 0,
+  newFileContent:
+`async function fn(a: Promise<string>, b: string) {
+  fn(a, await a);
+}`
+});

tests/cases/fourslash/codeFixAddMissingAwait_binaryExpressions.ts

@@ -0,0 +1,50 @@
+/// <reference path="fourslash.ts" />
+////async function fn(a: Promise<number>, b: number) {
+////  a | b;
+////  b + a;
+////  a + a;
+////}
+
+verify.codeFix({
+  description: ts.Diagnostics.Add_await.message,
+  index: 0,
+  newFileContent:
+`async function fn(a: Promise<number>, b: number) {
+  await a | b;
+  b + a;
+  a + a;
+}`
+});
+
+verify.codeFix({
+  description: ts.Diagnostics.Add_await.message,
+  index: 1,
+  newFileContent:
+`async function fn(a: Promise<number>, b: number) {
+  a | b;
+  b + await a;
+  a + a;
+}`
+});
+
+verify.codeFix({
+  description: ts.Diagnostics.Add_await.message,
+  index: 2,
+  newFileContent:
+`async function fn(a: Promise<number>, b: number) {
+  a | b;
+  b + a;
+  await a + await a;
+}`
+});
+
+verify.codeFixAll({
+  fixAllDescription: ts.Diagnostics.Fix_all_expressions_possibly_missing_await.message,
+  fixId: "addMissingAwait",
+  newFileContent:
+`async function fn(a: Promise<number>, b: number) {
+  await a | b;
+  b + await a;
+  await a + await a;
+}`
+});

tests/cases/fourslash/codeFixAddMissingAwait_initializer.ts

@@ -0,0 +1,28 @@
+/// <reference path="fourslash.ts" />
+////async function fn(a: string, b: Promise<string>) {
+////  const x = b;
+////  fn(x, b);
+////  fn(b, b);
+////}
+
+verify.codeFix({
+  description: "Add 'await' to initializer for 'x'",
+  index: 0,
+  newFileContent:
+`async function fn(a: string, b: Promise<string>) {
+  const x = await b;
+  fn(x, b);
+  fn(b, b);
+}`
+});
+
+verify.codeFixAll({
+  fixAllDescription: ts.Diagnostics.Fix_all_expressions_possibly_missing_await.message,
+  fixId: "addMissingAwait",
+  newFileContent:
+`async function fn(a: string, b: Promise<string>) {
+  const x = await b;
+  fn(x, b);
+  fn(await b, b);
+}`
+});

tests/cases/fourslash/codeFixAddMissingAwait_iterables.ts

@@ -0,0 +1,50 @@
+/// <reference path="fourslash.ts" />
+////async function fn(a: Promise<string[]>) {
+////  [...a];
+////  for (const c of a) { c; }
+////  for await (const c of a) { c; }
+////}
+
+verify.codeFix({
+  description: ts.Diagnostics.Add_await.message,
+  index: 0,
+  newFileContent:
+`async function fn(a: Promise<string[]>) {
+  [...await a];
+  for (const c of a) { c; }
+  for await (const c of a) { c; }
+}`
+});
+
+verify.codeFix({
+  description: ts.Diagnostics.Add_await.message,
+  index: 1,
+  newFileContent:
+`async function fn(a: Promise<string[]>) {
+  [...a];
+  for (const c of await a) { c; }
+  for await (const c of a) { c; }
+}`
+});
+
+verify.codeFix({
+  description: ts.Diagnostics.Add_await.message,
+  index: 2,
+  newFileContent:
+`async function fn(a: Promise<string[]>) {
+  [...a];
+  for (const c of a) { c; }
+  for await (const c of await a) { c; }
+}`
+});
+
+verify.codeFixAll({
+  fixAllDescription: ts.Diagnostics.Fix_all_expressions_possibly_missing_await.message,
+  fixId: "addMissingAwait",
+  newFileContent:
+`async function fn(a: Promise<string[]>) {
+  [...await a];
+  for (const c of await a) { c; }
+  for await (const c of await a) { c; }
+}`
+});

tests/cases/fourslash/codeFixAddMissingAwait_notAvailableWithoutPromise.ts

@@ -0,0 +1,6 @@
+/// <reference path="fourslash.ts" />
+////async function fn(a: {}, b: number) {
+////  a + b;
+////}
+
+verify.not.codeFixAvailable("addMissingAwait");

tests/cases/fourslash/codeFixAddMissingAwait_propertyAccess.ts

@@ -0,0 +1,13 @@
+/// <reference path="fourslash.ts" />
+////async function fn(a: Promise<{ x: string }>) {
+////  a.x;
+////}
+
+verify.codeFix({
+  description: ts.Diagnostics.Add_await.message,
+  index: 0,
+  newFileContent:
+`async function fn(a: Promise<{ x: string }>) {
+  (await a).x;
+}`
+});

tests/cases/fourslash/codeFixAddMissingAwait_signatures.ts

@@ -0,0 +1,50 @@
+/// <reference path="fourslash.ts" />
+////async function fn(a: Promise<() => void>, b: Promise<() => void> | (() => void), C: Promise<{ new(): any }>) {
+////  a();
+////  b();
+////  new C();
+////}
+
+verify.codeFix({
+  description: ts.Diagnostics.Add_await.message,
+  index: 0,
+  newFileContent:
+`async function fn(a: Promise<() => void>, b: Promise<() => void> | (() => void), C: Promise<{ new(): any }>) {
+  (await a)();
+  b();
+  new C();
+}`
+});
+
+verify.codeFix({
+  description: ts.Diagnostics.Add_await.message,
+  index: 1,
+  newFileContent:
+`async function fn(a: Promise<() => void>, b: Promise<() => void> | (() => void), C: Promise<{ new(): any }>) {
+  a();
+  (await b)();
+  new C();
+}`
+});
+
+verify.codeFix({
+  description: ts.Diagnostics.Add_await.message,
+  index: 2,
+  newFileContent:
+`async function fn(a: Promise<() => void>, b: Promise<() => void> | (() => void), C: Promise<{ new(): any }>) {
+  a();
+  b();
+  new (await C)();
+}`
+});
+
+verify.codeFixAll({
+  fixAllDescription: ts.Diagnostics.Fix_all_expressions_possibly_missing_await.message,
+  fixId: "addMissingAwait",
+  newFileContent:
+`async function fn(a: Promise<() => void>, b: Promise<() => void> | (() => void), C: Promise<{ new(): any }>) {
+  (await a)();
+  (await b)();
+  new (await C)();
+}`
+});

tests/cases/fourslash/codeFixAddMissingAwait_topLevel.ts

@@ -0,0 +1,10 @@
+/// <reference path="fourslash.ts" />
+////declare function getPromise(): Promise<string>;
+////const p = getPromise();
+////while (true) {
+////  p/*0*/.toLowerCase();
+////  getPromise()/*1*/.toLowerCase();
+////}
+
+verify.not.codeFixAvailable("addMissingAwait");
+verify.not.codeFixAvailable("addMissingAwaitToInitializer");

tests/cases/fourslash/fourslash.ts

@@ -185,7 +185,7 @@ declare namespace FourSlashInterface {
             applyChanges?: boolean,
             commands?: {}[],
         });
-        codeFixAvailable(options?: ReadonlyArray<VerifyCodeFixAvailableOptions>): void;
+        codeFixAvailable(options?: ReadonlyArray<VerifyCodeFixAvailableOptions> | string): void;
         applicableRefactorAvailableAtMarker(markerName: string): void;
         codeFixDiagnosticsAvailableAtMarkers(markerNames: string[], diagnosticCode?: number): void;
         applicableRefactorAvailableForRange(): void;