Improve asccuracy of remove unnecessary await fix

microsoft · Jul 12, 2019 · 3209373 · 3209373
1 parent 37f2e59
commit 3209373
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 4 deletions.
diff --git a/src/compiler/factory.ts b/src/compiler/factory.ts
@@ -4701,7 +4701,7 @@ namespace ts {
         }
     }
 
-    function getLeftmostExpression(node: Expression, stopAtCallExpressions: boolean) {
+    export function getLeftmostExpression(node: Expression, stopAtCallExpressions: boolean) {
         while (true) {
             switch (node.kind) {
                 case SyntaxKind.PostfixUnaryExpression:

diff --git a/src/services/codefixes/removeUnnecessaryAwait.ts b/src/services/codefixes/removeUnnecessaryAwait.ts
@@ -26,8 +26,18 @@ namespace ts.codefix {
             return;
         }
 
-        const parenthesizedExpression = tryCast(awaitExpression.parent, isParenthesizedExpression);
-        const removeParens = parenthesizedExpression && (isIdentifier(awaitExpression.expression) || isCallExpression(awaitExpression.expression));
-        changeTracker.replaceNode(sourceFile, removeParens ? parenthesizedExpression || awaitExpression : awaitExpression, awaitExpression.expression);
+        let expressionToReplace: Node = awaitExpression;
+        const hasSurroundingParens = isParenthesizedExpression(awaitExpression.parent);
+        if (hasSurroundingParens) {
+            const leftMostExpression = getLeftmostExpression(awaitExpression.expression, /*stopAtCallExpressions*/ false);
+            if (isIdentifier(leftMostExpression)) {
+                const precedingToken = findPrecedingToken(awaitExpression.parent.pos, sourceFile);
+                if (precedingToken && precedingToken.kind !== SyntaxKind.NewKeyword) {
+                    expressionToReplace = awaitExpression.parent;
+                }
+            }
+        }
+
+        changeTracker.replaceNode(sourceFile, expressionToReplace, awaitExpression.expression);
     }
 }
diff --git a/tests/cases/fourslash/codeFixRemoveUnnecessaryAwait.ts b/tests/cases/fourslash/codeFixRemoveUnnecessaryAwait.ts
@@ -1,26 +1,32 @@
 /// <reference path="fourslash.ts" />
 ////declare class C { foo(): void }
+////declare function getC(): { Class: C };
 ////declare function foo(): string;
 ////async function f() {
 ////  await "";
 ////  await 0;
 ////  (await foo()).toLowerCase();
 ////  (await 0).toFixed();
 ////  (await new C).foo();
+////  (await function() { }());
+////  new (await getC()).Class();
 ////}
 
 verify.codeFix({
   description: ts.Diagnostics.Remove_unnecessary_await.message,
   index: 0,
   newFileContent:
 `declare class C { foo(): void }
+declare function getC(): { Class: C };
 declare function foo(): string;
 async function f() {
   "";
   await 0;
   (await foo()).toLowerCase();
   (await 0).toFixed();
   (await new C).foo();
+  (await function() { }());
+  new (await getC()).Class();
 }`
 });
 
@@ -29,12 +35,15 @@ verify.codeFixAll({
   fixId: "removeUnnecessaryAwait",
   newFileContent:
 `declare class C { foo(): void }
+declare function getC(): { Class: C };
 declare function foo(): string;
 async function f() {
   "";
   0;
   foo().toLowerCase();
   (0).toFixed();
   (new C).foo();
+  (function() { } ());
+  new (getC()).Class();
 }`
 });