await-promise: check for-await-of (palantir#3297)

helenanderson · Oct 20, 2017 · 3881ddc · 3881ddc
1 parent db73e4b
commit 3881ddc
Show file tree

Hide file tree

Showing 6 changed files with 87 additions and 19 deletions.
diff --git a/src/rules/awaitPromiseRule.ts b/src/rules/awaitPromiseRule.ts
@@ -15,7 +15,7 @@
  * limitations under the License.
  */
 
-import { isAwaitExpression, isUnionOrIntersectionType } from "tsutils";
+import { isAwaitExpression, isForOfStatement, isTypeReference, isUnionOrIntersectionType } from "tsutils";
 import * as ts from "typescript";
 import * as Lint from "../index";
 
@@ -41,7 +41,8 @@ export class Rule extends Lint.Rules.TypedRule {
     };
     /* tslint:enable:object-literal-sort-keys */
 
-    public static FAILURE_STRING = "'await' of non-Promise.";
+    public static FAILURE_STRING = "Invalid 'await' of a non-Promise value.";
+    public static FAILURE_FOR_AWAIT_OF = "Invalid 'for-await-of' of a non-AsyncIterable value.";
 
     public applyWithProgram(sourceFile: ts.SourceFile, program: ts.Program): Lint.RuleFailure[] {
         const promiseTypes = new Set(["Promise", ...this.ruleArguments as string[]]);
@@ -54,27 +55,33 @@ function walk(ctx: Lint.WalkContext<Set<string>>, tc: ts.TypeChecker) {
     return ts.forEachChild(ctx.sourceFile, cb);
 
     function cb(node: ts.Node): void {
-        if (isAwaitExpression(node) && !couldBePromise(tc.getTypeAtLocation(node.expression))) {
+        if (isAwaitExpression(node) && !containsType(tc.getTypeAtLocation(node.expression), isPromiseType)) {
             ctx.addFailureAtNode(node, Rule.FAILURE_STRING);
+        } else if (isForOfStatement(node) && node.awaitModifier !== undefined &&
+                   !containsType(tc.getTypeAtLocation(node.expression), isAsyncIterable)) {
+            ctx.addFailureAtNode(node.expression, Rule.FAILURE_FOR_AWAIT_OF);
         }
         return ts.forEachChild(node, cb);
     }
 
-    function couldBePromise(type: ts.Type): boolean {
-        if (Lint.isTypeFlagSet(type, ts.TypeFlags.Any) || isPromiseType(type)) {
-            return true;
-        }
-
-        if (isUnionOrIntersectionType(type)) {
-            return type.types.some(couldBePromise);
-        }
-
-        const bases = type.getBaseTypes();
-        return bases !== undefined && bases.some(couldBePromise);
+    function isPromiseType(name: string) {
+        return promiseTypes.has(name);
     }
+}
 
-    function isPromiseType(type: ts.Type): boolean {
-        const { target } = type as ts.TypeReference;
-        return target !== undefined && target.symbol !== undefined && promiseTypes.has(target.symbol.name);
+function containsType(type: ts.Type, predicate: (name: string) => boolean): boolean {
+    if (Lint.isTypeFlagSet(type, ts.TypeFlags.Any) ||
+        isTypeReference(type) && type.target.symbol !== undefined && predicate(type.target.symbol.name)) {
+        return true;
     }
+    if (isUnionOrIntersectionType(type)) {
+        return type.types.some((t) => containsType(t, predicate));
+    }
+
+    const bases = type.getBaseTypes();
+    return bases !== undefined && bases.some((t) => containsType(t, predicate));
+}
+
+function isAsyncIterable(name: string) {
+    return name === "AsyncIterable" || name === "AsyncIterableIterator";
 }
diff --git a/test/rules/await-promise/custom-promise/test.ts.lint b/test/rules/await-promise/custom-promise/test.ts.lint
@@ -75,4 +75,4 @@ async function fCustomPromise() {
     await intersectionPromise;
 }
 
-[0]: 'await' of non-Promise.
+[0]: Invalid 'await' of a non-Promise value.
diff --git a/test/rules/await-promise/es6-promise/test.ts.lint b/test/rules/await-promise/es6-promise/test.ts.lint
@@ -48,4 +48,4 @@ async function fPromise() {
     await intersectionPromise;
 }
 
-[0]: 'await' of non-Promise.
+[0]: Invalid 'await' of a non-Promise value.
diff --git a/test/rules/await-promise/for-await-of/test.ts.lint b/test/rules/await-promise/for-await-of/test.ts.lint
@@ -0,0 +1,51 @@
+[typescript]: >= 2.3.0
+async function correct(foo: AsyncIterableIterator<string>) {
+    for await (const element of foo) {}
+}
+
+async function correct2() {
+    for await (const element of asyncGenerator()) {}
+}
+
+async function correct(foo: AsyncIterable<string>) {
+    for await (const element of foo) {}
+}
+
+async function correct3(foo: AsyncIterableIterator<string> | AsyncIterableIterator<number>) {
+    for await (const element of foo) {}
+}
+
+async function incorrect(foo: <Array<Promise<string>>) {
+    for await (const element of foo) {}
+                                ~~~ [0]
+}
+
+async function incorrect2(foo: IterableIterator<Promise<string>>) {
+    for await (const element of foo) {}
+                                ~~~ [0]
+}
+
+async function incorrect3() {
+    for await (const element of asyncGenerator) {}
+                                ~~~~~~~~~~~~~~ [0]
+}
+
+async function incorrect4() {
+    for await (const element of generator()) {}
+                                ~~~~~~~~~~~ [0]
+}
+
+async function incorrect5(foo: Iterable<string>) {
+    for await (const element of foo) {}
+                                ~~~ [0]
+}
+
+async function* asyncGenerator() {
+    yield 1;
+}
+
+function* generator() {
+    yield 1;
+}
+
+[0]: Invalid 'for-await-of' of a non-AsyncIterable value.
diff --git a/test/rules/await-promise/for-await-of/tsconfig.json b/test/rules/await-promise/for-await-of/tsconfig.json
@@ -0,0 +1,5 @@
+{
+    "compilerOptions": {
+        "target": "esnext"
+    }
+}
diff --git a/test/rules/await-promise/for-await-of/tslint.json b/test/rules/await-promise/for-await-of/tslint.json
@@ -0,0 +1,5 @@
+{
+  "rules": {
+    "await-promise": true
+  }
+}