no-array-for-each: Add fix for arrow function body (#1785)

sindresorhus · Apr 6, 2022 · 59218e3 · 59218e3
1 parent 0034e69
commit 59218e3
Show file tree

Hide file tree

Showing 6 changed files with 231 additions and 8 deletions.
diff --git a/rules/ast/index.js b/rules/ast/index.js
@@ -1,4 +1,5 @@
 'use strict';
 module.exports = {
+	isArrowFunctionBody: require('./is-arrow-function-body.js'),
 	isEmptyNode: require('./is-empty-node.js'),
 };
diff --git a/rules/ast/is-arrow-function-body.js b/rules/ast/is-arrow-function-body.js
@@ -0,0 +1,7 @@
+'use strict';
+
+function isArrowFunctionBody(node) {
+	return node.parent.type === 'ArrowFunctionExpression' && node.parent.body === node;
+}
+
+module.exports = isArrowFunctionBody;
diff --git a/rules/no-array-for-each.js b/rules/no-array-for-each.js
@@ -18,6 +18,7 @@ const isFunctionSelfUsedInside = require('./utils/is-function-self-used-inside.j
 const {isNodeMatches} = require('./utils/is-node-matches.js');
 const assertToken = require('./utils/assert-token.js');
 const {fixSpaceAroundKeyword, removeParentheses} = require('./fix/index.js');
+const {isArrowFunctionBody} = require('./ast/index.js');
 
 const MESSAGE_ID_ERROR = 'no-array-for-each/error';
 const MESSAGE_ID_SUGGESTION = 'no-array-for-each/suggestion';
@@ -83,7 +84,7 @@ function getFixFunction(callExpression, functionInfo, context) {
 	const iterableObject = callExpression.callee.object;
 	const {returnStatements} = functionInfo.get(callback);
 	const isOptionalObject = callExpression.callee.optional;
-	const expressionStatement = stripChainExpression(callExpression).parent;
+	const ancestor = stripChainExpression(callExpression).parent;
 	const objectText = sourceCode.getText(iterableObject);
 
 	const getForOfLoopHeadText = () => {
@@ -247,12 +248,17 @@ function getFixFunction(callExpression, functionInfo, context) {
 			yield * replaceReturnStatement(returnStatement, fixer);
 		}
 
-		const expressionStatementLastToken = sourceCode.getLastToken(expressionStatement);
-		// Remove semicolon if it's not needed anymore
-		// foo.forEach(bar => {});
-		//                       ^
-		if (shouldRemoveExpressionStatementLastToken(expressionStatementLastToken)) {
-			yield fixer.remove(expressionStatementLastToken, fixer);
+		if (ancestor.type === 'ExpressionStatement') {
+			const expressionStatementLastToken = sourceCode.getLastToken(ancestor);
+			// Remove semicolon if it's not needed anymore
+			// foo.forEach(bar => {});
+			//                       ^
+			if (shouldRemoveExpressionStatementLastToken(expressionStatementLastToken)) {
+				yield fixer.remove(expressionStatementLastToken, fixer);
+			}
+		} else if (ancestor.type === 'ArrowFunctionExpression') {
+			yield fixer.insertTextBefore(callExpression, '{ ');
+			yield fixer.insertTextAfter(callExpression, ' }');
 		}
 
 		yield * fixSpaceAroundKeyword(fixer, callExpression.parent, sourceCode);
@@ -332,7 +338,11 @@ function isFixable(callExpression, {scope, functionInfo, allIdentifiers, context
 	}
 
 	// Check ancestors, we only fix `ExpressionStatement`
-	if (stripChainExpression(callExpression).parent.type !== 'ExpressionStatement') {
+	const callOrChainExpression = stripChainExpression(callExpression);
+	if (
+		callOrChainExpression.parent.type !== 'ExpressionStatement'
+		&& !isArrowFunctionBody(callOrChainExpression)
+	) {
 		return false;
 	}
 

diff --git a/test/no-array-for-each.mjs b/test/no-array-for-each.mjs
@@ -425,10 +425,21 @@ test.snapshot({
 
 			// Need insert space after keyword
 			'if (true) {} else[foo].forEach((element) => {})',
+
+			// Arrow function body
+			'const a = () => ((  foo.forEach(element => bar(element))  ))',
+			'const a = () => ((  foo.forEach(element => bar(element))  ));',
+			'const a = () => foo.forEach(element => bar(element))',
+			'const a = () => foo.forEach(element => bar(element));',
+			'const a = () => void foo.forEach(element => bar(element));',
 		].flatMap(code => [code, code.replace('.forEach', '?.forEach')]),
 
 		// Should not fix to invalid code
 		'1?.forEach((a, b) => call(a, b))',
+
+		// Arrow function body
+		'array.forEach((arrayInArray) => arrayInArray.forEach(element => bar(element)));',
+		'array.forEach((arrayInArray) => arrayInArray?.forEach(element => bar(element)));',
 	],
 });
 

diff --git a/test/snapshots/no-array-for-each.mjs.md b/test/snapshots/no-array-for-each.mjs.md
@@ -3897,6 +3897,154 @@ Generated by [AVA](https://avajs.dev).
     `
 
 ## Invalid #216
+      1 | const a = () => ((  foo.forEach(element => bar(element))  ))
+
+> Output
+
+    `␊
+      1 | const a = () =>   { for (const element of foo)  bar(element) }  ␊
+    `
+
+> Error 1/1
+
+    `␊
+    > 1 | const a = () => ((  foo.forEach(element => bar(element))  ))␊
+        |                         ^^^^^^^ Use \`for…of\` instead of \`.forEach(…)\`.␊
+    `
+
+## Invalid #217
+      1 | const a = () => ((  foo?.forEach(element => bar(element))  ))
+
+> Output
+
+    `␊
+      1 | const a = () =>   { if (foo) for (const element of foo)  bar(element) }  ␊
+    `
+
+> Error 1/1
+
+    `␊
+    > 1 | const a = () => ((  foo?.forEach(element => bar(element))  ))␊
+        |                          ^^^^^^^ Use \`for…of\` instead of \`.forEach(…)\`.␊
+    `
+
+## Invalid #218
+      1 | const a = () => ((  foo.forEach(element => bar(element))  ));
+
+> Output
+
+    `␊
+      1 | const a = () =>   { for (const element of foo)  bar(element) }  ;␊
+    `
+
+> Error 1/1
+
+    `␊
+    > 1 | const a = () => ((  foo.forEach(element => bar(element))  ));␊
+        |                         ^^^^^^^ Use \`for…of\` instead of \`.forEach(…)\`.␊
+    `
+
+## Invalid #219
+      1 | const a = () => ((  foo?.forEach(element => bar(element))  ));
+
+> Output
+
+    `␊
+      1 | const a = () =>   { if (foo) for (const element of foo)  bar(element) }  ;␊
+    `
+
+> Error 1/1
+
+    `␊
+    > 1 | const a = () => ((  foo?.forEach(element => bar(element))  ));␊
+        |                          ^^^^^^^ Use \`for…of\` instead of \`.forEach(…)\`.␊
+    `
+
+## Invalid #220
+      1 | const a = () => foo.forEach(element => bar(element))
+
+> Output
+
+    `␊
+      1 | const a = () => { for (const element of foo)  bar(element) }␊
+    `
+
+> Error 1/1
+
+    `␊
+    > 1 | const a = () => foo.forEach(element => bar(element))␊
+        |                     ^^^^^^^ Use \`for…of\` instead of \`.forEach(…)\`.␊
+    `
+
+## Invalid #221
+      1 | const a = () => foo?.forEach(element => bar(element))
+
+> Output
+
+    `␊
+      1 | const a = () => { if (foo) for (const element of foo)  bar(element) }␊
+    `
+
+> Error 1/1
+
+    `␊
+    > 1 | const a = () => foo?.forEach(element => bar(element))␊
+        |                      ^^^^^^^ Use \`for…of\` instead of \`.forEach(…)\`.␊
+    `
+
+## Invalid #222
+      1 | const a = () => foo.forEach(element => bar(element));
+
+> Output
+
+    `␊
+      1 | const a = () => { for (const element of foo)  bar(element) };␊
+    `
+
+> Error 1/1
+
+    `␊
+    > 1 | const a = () => foo.forEach(element => bar(element));␊
+        |                     ^^^^^^^ Use \`for…of\` instead of \`.forEach(…)\`.␊
+    `
+
+## Invalid #223
+      1 | const a = () => foo?.forEach(element => bar(element));
+
+> Output
+
+    `␊
+      1 | const a = () => { if (foo) for (const element of foo)  bar(element) };␊
+    `
+
+> Error 1/1
+
+    `␊
+    > 1 | const a = () => foo?.forEach(element => bar(element));␊
+        |                      ^^^^^^^ Use \`for…of\` instead of \`.forEach(…)\`.␊
+    `
+
+## Invalid #224
+      1 | const a = () => void foo.forEach(element => bar(element));
+
+> Error 1/1
+
+    `␊
+    > 1 | const a = () => void foo.forEach(element => bar(element));␊
+        |                          ^^^^^^^ Use \`for…of\` instead of \`.forEach(…)\`.␊
+    `
+
+## Invalid #225
+      1 | const a = () => void foo?.forEach(element => bar(element));
+
+> Error 1/1
+
+    `␊
+    > 1 | const a = () => void foo?.forEach(element => bar(element));␊
+        |                           ^^^^^^^ Use \`for…of\` instead of \`.forEach(…)\`.␊
+    `
+
+## Invalid #226
       1 | 1?.forEach((a, b) => call(a, b))
 
 > Output
@@ -3911,3 +4059,49 @@ Generated by [AVA](https://avajs.dev).
     > 1 | 1?.forEach((a, b) => call(a, b))␊
         |    ^^^^^^^ Use \`for…of\` instead of \`.forEach(…)\`.␊
     `
+
+## Invalid #227
+      1 | array.forEach((arrayInArray) => arrayInArray.forEach(element => bar(element)));
+
+> Output
+
+    `␊
+      1 | for (const arrayInArray of array)  for (const element of arrayInArray)  bar(element);␊
+    `
+
+> Error 1/2
+
+    `␊
+    > 1 | array.forEach((arrayInArray) => arrayInArray.forEach(element => bar(element)));␊
+        |       ^^^^^^^ Use \`for…of\` instead of \`.forEach(…)\`.␊
+    `
+
+> Error 2/2
+
+    `␊
+    > 1 | array.forEach((arrayInArray) => arrayInArray.forEach(element => bar(element)));␊
+        |                                              ^^^^^^^ Use \`for…of\` instead of \`.forEach(…)\`.␊
+    `
+
+## Invalid #228
+      1 | array.forEach((arrayInArray) => arrayInArray?.forEach(element => bar(element)));
+
+> Output
+
+    `␊
+      1 | for (const arrayInArray of array)  if (arrayInArray) for (const element of arrayInArray)  bar(element);␊
+    `
+
+> Error 1/2
+
+    `␊
+    > 1 | array.forEach((arrayInArray) => arrayInArray?.forEach(element => bar(element)));␊
+        |       ^^^^^^^ Use \`for…of\` instead of \`.forEach(…)\`.␊
+    `
+
+> Error 2/2
+
+    `␊
+    > 1 | array.forEach((arrayInArray) => arrayInArray?.forEach(element => bar(element)));␊
+        |                                               ^^^^^^^ Use \`for…of\` instead of \`.forEach(…)\`.␊
+    `
diff --git a/test/snapshots/no-array-for-each.mjs.snap b/test/snapshots/no-array-for-each.mjs.snap