Await promise configurable (palantir#2705)

src/rules/awaitPromiseRule.ts

@@ -15,6 +15,7 @@
  * limitations under the License.
  */
 
+import { isAwaitExpression } from "tsutils";
 import * as ts from "typescript";
 import * as Lint from "../index";
 
@@ -23,9 +24,17 @@ export class Rule extends Lint.Rules.TypedRule {
     public static metadata: Lint.IRuleMetadata = {
         ruleName: "await-promise",
         description: "Warns for an awaited value that is not a Promise.",
-        optionsDescription: "Not configurable.",
-        options: null,
-        optionExamples: [true],
+        optionsDescription: Lint.Utils.dedent`
+            A list of 'string' names of any additional classes that should also be handled as Promises.
+        `,
+        options: {
+            type: "list",
+            listType: {
+                type: "array",
+                items: {type: "string"},
+            },
+        },
+        optionExamples: [true, [true, "Thenable"]],
         type: "functionality",
         typescriptOnly: true,
         requiresTypeInfo: true,
@@ -34,38 +43,40 @@ export class Rule extends Lint.Rules.TypedRule {
 
     public static FAILURE_STRING = "'await' of non-Promise.";
 
-    public applyWithProgram(srcFile: ts.SourceFile, program: ts.Program): Lint.RuleFailure[] {
-        return this.applyWithWalker(new Walker(srcFile, this.getOptions(), program));
+    public applyWithProgram(sourceFile: ts.SourceFile, program: ts.Program): Lint.RuleFailure[] {
+        const promiseTypes = new Set(["Promise", ...this.ruleArguments as string[]]);
+        const tc = program.getTypeChecker();
+        return this.applyWithFunction(sourceFile, (ctx) => walk(ctx, tc, promiseTypes));
     }
 }
 
-class Walker extends Lint.ProgramAwareRuleWalker {
-    public visitNode(node: ts.Node) {
-        if (node.kind === ts.SyntaxKind.AwaitExpression &&
-            !couldBePromise(this.getTypeChecker().getTypeAtLocation((node as ts.AwaitExpression).expression))) {
-            this.addFailureAtNode(node, Rule.FAILURE_STRING);
-        }
+function walk(ctx: Lint.WalkContext<void>, tc: ts.TypeChecker, promiseTypes: Set<string>) {
+    return ts.forEachChild(ctx.sourceFile, cb);
 
-        super.visitNode(node);
+    function cb(node: ts.Node): void {
+        if (isAwaitExpression(node) && !couldBePromise(tc.getTypeAtLocation(node.expression))) {
+            ctx.addFailureAtNode(node, Rule.FAILURE_STRING);
+        }
+        return ts.forEachChild(node, cb);
     }
-}
 
-function couldBePromise(type: ts.Type): boolean {
-    if (Lint.isTypeFlagSet(type, ts.TypeFlags.Any) || isPromiseType(type)) {
-        return true;
-    }
+    function couldBePromise(type: ts.Type): boolean {
+        if (Lint.isTypeFlagSet(type, ts.TypeFlags.Any) || isPromiseType(type)) {
+            return true;
+        }
 
-    if (isUnionType(type)) {
-        return type.types.some(isPromiseType);
-    }
+        if (isUnionType(type)) {
+            return type.types.some(couldBePromise);
+        }
 
-    const bases = type.getBaseTypes();
-    return bases !== undefined && bases.some(couldBePromise);
-}
+        const bases = type.getBaseTypes();
+        return bases !== undefined && bases.some(couldBePromise);
+    }
 
-function isPromiseType(type: ts.Type): boolean {
-    const { target } = type as ts.TypeReference;
-    return target !== undefined && target.symbol !== undefined && target.symbol.name === "Promise";
+    function isPromiseType(type: ts.Type): boolean {
+        const { target } = type as ts.TypeReference;
+        return target !== undefined && target.symbol !== undefined && promiseTypes.has(target.symbol.name);
+    }
 }
 
 function isUnionType(type: ts.Type): type is ts.UnionType {

test/rules/await-promise/custom-promise/test.ts.lint

@@ -0,0 +1,70 @@
+class CustomPromise<T> {}
+
+async function fAny() {
+    const isAny: any = 1;
+
+    // any type
+    await isAny;
+
+    // union type with any type
+    await (Math.random() > 0.5 ? isAny : 0);
+}
+
+async function fNonPromise() {
+    // number type
+    await 0;
+    ~~~~~~~ [0]
+
+    // union type without Promise
+    await (Math.random() > 0.5 ? "" : 0);
+    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ [0]
+
+    // non Promise type
+    class NonPromise extends Array {}
+    await new NonPromise();
+    ~~~~~~~~~~~~~~~~~~~~~~ [0]
+}
+
+async function fStandardPromise() {
+    // direct type
+    const numberPromise: Promise<number>;
+    await numberPromise;
+
+    // 1st level base type
+    class Foo extends Promise<number> {}
+    const foo: Foo = Foo.resolve(2);
+    await foo;
+
+    // 2nd level base type
+    class Bar extends Foo {}
+    const bar: Bar = Bar.resolve(2);
+    await bar;
+
+    // union type
+    await (Math.random() > 0.5 ? numberPromise : 0);
+    await (Math.random() > 0.5 ? foo : 0);
+    await (Math.random() > 0.5 ? bar : 0);
+}
+
+async function fCustomPromise() {
+    // direct type
+    const numberPromise: CustomPromise<number>;
+    await numberPromise;
+
+    // 1st level base type
+    class Foo extends CustomPromise<number> {}
+    const foo: Foo = Foo.resolve(2);
+    await foo;
+
+    // 2nd level base type
+    class Bar extends Foo {}
+    const bar: Bar = Bar.resolve(2);
+    await bar;
+
+    // union type
+    await (Math.random() > 0.5 ? numberPromise : 0);
+    await (Math.random() > 0.5 ? foo : 0);
+    await (Math.random() > 0.5 ? bar : 0);
+}
+
+[0]: 'await' of non-Promise.

test/rules/await-promise/custom-promise/tslint.json

@@ -0,0 +1,8 @@
+{
+  "linterOptions": {
+    "typeCheck": true
+  },
+  "rules": {
+    "await-promise": [true, "CustomPromise"]
+  }
+}

test/rules/await-promise/es6-promise/test.ts.lint

@@ -0,0 +1,47 @@
+async function fAny() {
+    const isAny: any = 1;
+
+    // any type
+    await isAny;
+
+    // union type with any type
+    await (Math.random() > 0.5 ? isAny : 0);
+}
+
+async function fNonPromise() {
+    // number type
+    await 0;
+    ~~~~~~~ [0]
+
+    // union type without Promise
+    await (Math.random() > 0.5 ? "" : 0);
+    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ [0]
+
+    // non Promise type
+    class NonPromise extends Array {}
+    await new NonPromise();
+    ~~~~~~~~~~~~~~~~~~~~~~ [0]
+}
+
+async function fPromise() {
+    // direct type
+    const numberPromise: Promise<number>;
+    await numberPromise;
+
+    // 1st level base type
+    class Foo extends Promise<number> {}
+    const foo: Foo = Foo.resolve(2);
+    await foo;
+
+    // 2nd level base type
+    class Bar extends Foo {}
+    const bar: Bar = Bar.resolve(2);
+    await bar;
+
+    // union type
+    await (Math.random() > 0.5 ? numberPromise : 0);
+    await (Math.random() > 0.5 ? foo : 0);
+    await (Math.random() > 0.5 ? bar : 0);
+}
+
+[0]: 'await' of non-Promise.

test/rules/await-promise/es6-promise/tsconfig.json

@@ -0,0 +1,5 @@
+{
+    "compilerOptions": {
+        "target": "es6"
+    }
+}