Require code to be inside a function body to offer await

Author: andrewbranch

src/harness/fourslash.ts

@@ -2828,11 +2828,28 @@ Actual: ${stringify(fullActual)}`);
             }
         }
 
-        public verifyCodeFixAvailable(negative: boolean, expected: FourSlashInterface.VerifyCodeFixAvailableOptions[] | undefined): void {
-            assert(!(negative && expected), "Cannot provide an expected code fix with a negated assertion");
+        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(({ fixName }) => 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

@@ -88,14 +88,21 @@ namespace ts.codefix {
 
     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 isMissingAwaitError(sourceFile, errorCode, span, cancellationToken, program) ? expression : undefined;
+        return expression
+            && isMissingAwaitError(sourceFile, errorCode, span, cancellationToken, program)
+            && isInsideAwaitableBody(expression)
+                ? expression
+                : undefined;
     }
 
     function findAwaitableInitializer(expression: Node, sourceFile: SourceFile, checker: TypeChecker): Expression | undefined {
@@ -116,7 +123,8 @@ namespace ts.codefix {
             !declaration.initializer ||
             variableStatement.getSourceFile() !== sourceFile ||
             hasModifier(variableStatement, ModifierFlags.Export) ||
-            !variableName) {
+            !variableName ||
+            !isInsideAwaitableBody(declaration.initializer)) {
             return;
         }
 
@@ -131,6 +139,16 @@ namespace ts.codefix {
         return declaration.initializer;
     }
 
+    function isInsideAwaitableBody(node: Node) {
+        return !!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;

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;