prefer-top-level-await: Improve top-level expression detection (#1526)

Author: fisker

rules/prefer-top-level-await.js

@@ -13,7 +13,9 @@ const messages = {
 	[SUGGESTION_ADD_AWAIT]: 'Insert `await`.',
 };
 
-const topLevelCallExpression = 'Program > ExpressionStatement > CallExpression[optional!=true].expression';
+const promiseMethods = ['then', 'catch', 'finally'];
+
+const topLevelCallExpression = 'CallExpression:not(:function *)';
 const iife = [
 	topLevelCallExpression,
 	matches([
@@ -27,31 +29,60 @@ const promise = [
 	topLevelCallExpression,
 	memberExpressionSelector({
 		path: 'callee',
-		properties: ['then', 'catch', 'finally'],
+		properties: promiseMethods,
+		includeOptional: true,
 	}),
 ].join('');
 const identifier = [
 	topLevelCallExpression,
 	'[callee.type="Identifier"]',
 ].join('');
 
+const isPromiseMethodCalleeObject = node =>
+	node.parent.type === 'MemberExpression'
+	&& node.parent.object === node
+	&& !node.parent.computed
+	&& node.parent.property.type === 'Identifier'
+	&& promiseMethods.includes(node.parent.property.name)
+	&& node.parent.parent.type === 'CallExpression'
+	&& node.parent.parent.callee === node.parent;
+const isAwaitArgument = node => {
+	if (node.parent.type === 'ChainExpression') {
+		node = node.parent;
+	}
+
+	return node.parent.type === 'AwaitExpression' && node.parent.argument === node;
+};
+
 /** @param {import('eslint').Rule.RuleContext} context */
 function create(context) {
 	return {
 		[promise](node) {
+			if (isPromiseMethodCalleeObject(node) || isAwaitArgument(node)) {
+				return;
+			}
+
 			return {
 				node: node.callee.property,
 				messageId: ERROR_PROMISE,
 			};
 		},
 		[iife](node) {
+			if (isPromiseMethodCalleeObject(node) || isAwaitArgument(node)) {
+				return;
+			}
+
 			return {
 				node,
 				loc: getFunctionHeadLocation(node.callee, context.getSourceCode()),
 				messageId: ERROR_IIFE,
 			};
 		},
 		[identifier](node) {
+			if (isPromiseMethodCalleeObject(node) || isAwaitArgument(node)) {
+				return;
+			}
+
 			const variable = findVariable(context.getScope(), node.callee);
 			if (!variable || variable.defs.length !== 1) {
 				return;

test/prefer-top-level-await.mjs

@@ -8,45 +8,64 @@ test.snapshot({
 	valid: [
 		'a()',
 		'a = async () => {}',
-		'a = (async () => {})()',
+		'(async function *() {})()',
 		outdent`
-			{
-				(async () => {})();
+			function foo() {
+				if (foo) {
+					(async () => {})()
+				}
 			}
 		`,
-		'!async function() {}()',
-		'void async function() {}()',
-		'(async function *() {})()',
-		'(async () => {})?.()',
+		'await (async () => {})()',
 	],
 	invalid: [
 		'(async () => {})()',
+		'(async () => {})?.()',
 		'(async function() {})()',
 		'(async function() {}())',
 		'(async function run() {})()',
 		'(async function(c, d) {})(a, b)',
+		'if (foo) (async () => {})()',
+		outdent`
+			{
+				(async () => {})();
+			}
+		`,
+		'a = (async () => {})()',
+		'!async function() {}()',
+		'void async function() {}()',
+		'(async () => {})().catch(foo)',
 	],
 });
 
 // Promise
 test.snapshot({
 	valid: [
 		'foo.then',
-		'foo.then().toString()',
-		'!foo.then()',
-		'foo.then?.(bar)',
-		'foo?.then(bar)',
-		'foo?.then(bar).finally(qux)',
+		'await foo.then(bar)',
+		'await foo.then(bar).catch(bar)',
+		'await foo.then?.(bar)',
+		'await foo.then(bar)?.catch(bar)',
+		'await foo.then(bar)?.catch?.(bar)',
 	],
 	invalid: [
 		'foo.then(bar)',
+		'foo.then?.(bar)',
+		'foo?.then(bar)',
 		'foo.catch(() => process.exit(1))',
 		'foo.finally(bar)',
 		'foo.then(bar, baz)',
 		'foo.then(bar, baz).finally(qux)',
 		'(foo.then(bar, baz)).finally(qux)',
 		'(async () => {})().catch(() => process.exit(1))',
 		'(async function() {}()).finally(() => {})',
+		'for (const foo of bar) foo.then(bar)',
+		'foo?.then(bar).finally(qux)',
+		'foo.then().toString()',
+		'!foo.then()',
+		'foo.then(bar).then(baz)?.then(qux)',
+		'foo.then(bar).then(baz).then?.(qux)',
+		'foo.then(bar).catch(bar).finally(bar)',
 	],
 });
 
@@ -106,10 +125,6 @@ test.snapshot({
 			`,
 			parserOptions: {sourceType: 'script'},
 		},
-		outdent`
-			const foo = async () => {};
-			foo?.();
-		`,
 		outdent`
 			const program = {async run () {}};
 			program.run()
@@ -119,12 +134,24 @@ test.snapshot({
 			const {run} = program;
 			run()
 		`,
+		outdent`
+			const foo = async () => {};
+			await foo();
+		`,
 	],
 	invalid: [
 		outdent`
 			const foo = async () => {};
 			foo();
 		`,
+		outdent`
+			const foo = async () => {};
+			foo?.();
+		`,
+		outdent`
+			const foo = async () => {};
+			foo().then(foo);
+		`,
 		outdent`
 			const foo = async function () {}, bar = 1;
 			foo(bar);
@@ -133,6 +160,14 @@ test.snapshot({
 			foo();
 			async function foo() {}
 		`,
+		outdent`
+			const foo = async () => {};
+			if (true) {
+				alert();
+			} else {
+				foo();
+			}
+		`,
 	],
 });