fix: not properly fixing hooks return statements

arthurgeron · Dec 4, 2023 · c8e78a9 · c8e78a9
1 parent a657d0c
commit c8e78a9
Show file tree

Hide file tree

Showing 2 changed files with 78 additions and 3 deletions.
diff --git a/__tests__/require-usememo.test.ts b/__tests__/require-usememo.test.ts
@@ -560,7 +560,6 @@ describe('Rule - Require-usememo', () =>  {
         }`,
         errors: [{ messageId: "object-usememo-props" }],
       },
-
       {
         code: `
         const useTest = () => {
@@ -651,6 +650,62 @@ import type { ComponentProps } from 'react';
         }`,
         errors: [{ messageId: "array-usememo-props" }],
       },
+      {
+        code: `
+        function useTest() {
+          return { x: 1 };
+        }`,
+        output: `import { useMemo } from 'react';
+
+        function useTest() {
+          return useMemo(() => ({ x: 1 }), []);
+        }`,
+        errors: [{ messageId: "object-usememo-hook" }],
+        options: [{ checkHookReturnObject: true, strict: true }],
+      },
+      {
+        code: `
+        function useTest() {
+          return (data) => {};
+        }`,
+        output: `import { useCallback } from 'react';
+
+        function useTest() {
+          return useCallback((data) => {}, []);
+        }`,
+        errors: [{ messageId: "function-usecallback-hook" }],
+        options: [{ checkHookReturnObject: true, strict: true }],
+      },
+      {
+        code: `
+        function useTest() {
+          const y = { x: 1 };
+          return y;
+        }`,
+        output: `import { useMemo } from 'react';
+
+        function useTest() {
+          const y = useMemo(() => ({ x: 1 }), []);
+          return y;
+        }`,
+        errors: [{ messageId: "object-usememo-hook" }],
+        options: [{ checkHookReturnObject: true, strict: true }],
+      },
+      {
+        code: `
+        function useTest() {
+          const y = (data) => {};
+          return y;
+        }`,
+        output: `import { useCallback } from 'react';
+
+        function useTest() {
+          const y = useCallback((data) => {}, []);
+          return y;
+        }`,
+        errors: [{ messageId: "function-usecallback-hook" }],
+        options: [{ checkHookReturnObject: true, strict: true }],
+      },
     ],
   });
 });

diff --git a/src/require-usememo/utils.ts b/src/require-usememo/utils.ts
@@ -175,11 +175,11 @@ function getSafeVariableName(context: Rule.RuleContext, name: string, attempts =
 }
 
 // Eslint Auto-fix logic, functional components/hooks only
-export function fixBasedOnMessageId(node: Rule.Node | TSESTree.JSXElement, messageId: keyof typeof MessagesRequireUseMemo, fixer: Rule.RuleFixer, context: Rule.RuleContext, reactImportData: ReactImportInformation) {
+export function fixBasedOnMessageId(node: Rule.Node, messageId: keyof typeof MessagesRequireUseMemo, fixer: Rule.RuleFixer, context: Rule.RuleContext, reactImportData: ReactImportInformation) {
   const sourceCode = context.getSourceCode();
   let hook = messageIdToHookDict[messageId] || 'useMemo';
   const isObjExpression = node.type === 'ObjectExpression';
-  const isJSXElement = node.type === 'JSXElement';
+  const isJSXElement = (node as unknown as TSESTree.JSXElement).type === 'JSXElement';
   const parentIsVariableDeclarator = (node as Rule.Node).parent.type === 'VariableDeclarator';
   const isArrowFunctionExpression = node.type === 'ArrowFunctionExpression';
   const isFunctionExpression = node.type === 'FunctionExpression';
@@ -188,6 +188,26 @@ export function fixBasedOnMessageId(node: Rule.Node | TSESTree.JSXElement, messa
 
   // Determine what type of behavior to follow according to the error message
   switch (messageId) {
+    case 'function-usecallback-hook':
+      if (node.type === 'FunctionExpression' || node.type === 'ArrowFunctionExpression') {
+        const importStatementFixes = addReactImports(context, 'useCallback', reactImportData, fixer);
+        const fixed = fixFunction(node as TSESTree.FunctionExpression, context);
+        importStatementFixes && fixes.push(importStatementFixes);
+        fixes.push(fixer.replaceText(node as Rule.Node, fixed));
+        return fixes;
+      }
+      break;
+    case 'object-usememo-hook':
+      const _returnNode = node as TSESTree.ReturnStatement;
+      // An undefined node.argument means returned value is not an expression, but most probably a variable which should not be handled here, which falls under default, simpler fix logic.
+      if(_returnNode.argument) {
+        const importStatementFixes = addReactImports(context, 'useMemo', reactImportData, fixer);
+        const fixed = `useMemo(() => (${sourceCode.getText(_returnNode.argument as Rule.Node)}), [])`;
+        importStatementFixes && fixes.push(importStatementFixes);
+        fixes.push(fixer.replaceText(_returnNode.argument as Rule.Node, fixed));
+        return fixes;
+      } 
+      break;
     case 'function-usecallback-props':
     case 'object-usememo-props':
     case 'jsx-usememo-props':