add support for components wrapped in forwardRef

Author: kanzelm3

.vscode/launch.json

@@ -0,0 +1,17 @@
+{
+  "version": "0.2.0",
+  "configurations": [
+    {
+      "name": "Debug Jest Tests",
+      "type": "node",
+      "request": "launch",
+      "runtimeArgs": [
+        "--inspect-brk",
+        "${workspaceRoot}/node_modules/.bin/jest",
+        "--runInBand"
+      ],
+      "console": "integratedTerminal",
+      "internalConsoleOptions": "neverOpen"
+    }
+  ]
+}

index.js

@@ -36,15 +36,29 @@ function getReturnStatement(node) {
     return;
   }
 
+  if (node.type === 'ArrowFunctionExpression') return node.body;
   return node.type === 'VariableDeclaration'
     ? node.declarations?.[0]?.init?.body?.body?.find(
         (statement) => statement.type === 'ReturnStatement',
-      ) ?? node.declarations?.[0]?.init?.body
+      ) ??
+        node.declarations?.[0]?.init?.arguments?.[0]?.body ??
+        node.declarations?.[0]?.init?.body
     : node.body?.body?.find(
         (statement) => statement.type === 'ReturnStatement',
       );
 }
 
+function isForwardRef(node) {
+  if (!Boolean(node)) {
+    return;
+  }
+
+  return node.type === 'VariableDeclaration'
+    ? node.declarations?.[0]?.init?.callee?.property?.name === 'forwardRef'
+    : node.callee?.name === 'forwardRef' ||
+        node.callee?.property?.name === 'forwardRef';
+}
+
 function isTreeDone(node, excludeComponentNames) {
   return (
     node.type === 'JSXElement' &&
@@ -95,11 +109,20 @@ const rules = {
       return {
         Program(node) {
           const componentNodes = node.body
-            .map((child) => child?.declaration ?? child)
+            .map((child) => {
+              const declaration = child?.declaration ?? child;
+              if (isForwardRef(declaration)) {
+                // do something
+                return declaration?.arguments?.[0] ?? declaration;
+              }
+              return declaration;
+            })
             .filter(
               (child) =>
                 child.type === 'VariableDeclaration' ||
-                child.type === 'FunctionDeclaration',
+                child.type === 'FunctionDeclaration' ||
+                child.type === 'FunctionExpression' ||
+                child.type === 'ArrowFunctionExpression',
             )
             .filter((child) => {
               let flag = false;

test.js

@@ -91,8 +91,7 @@ const fragmentsWontUpdate = `const Component = () => {
 const reactForwardRef = /* tsx */ `
 export const InternalLink = React.forwardRef<HTMLAnchorElement, InternalLinkProps>(
   ({ variant, ...props }, ref) => (
-    <Link
-      data-component="InternalLink"
+    <Link data-component="InternalLink"
       ref={ref}
       className={classNames(css.link, { [css.inverse]: variant === 'inverse' })}
       {...props}
@@ -118,8 +117,7 @@ export const InternalLink = React.forwardRef<HTMLAnchorElement, InternalLinkProp
 const forwardRef = /* tsx */ `
 export const InternalLink = forwardRef<HTMLAnchorElement, InternalLinkProps>(
   ({ variant, ...props }, ref) => (
-    <Link
-      data-component="InternalLink"
+    <Link data-component="InternalLink"
       ref={ref}
       className={classNames(css.link, { [css.inverse]: variant === 'inverse' })}
       {...props}
@@ -142,6 +140,66 @@ export const InternalLink = forwardRef<HTMLAnchorElement, InternalLinkProps>(
   ),
 );`;
 
+const defaultReactForwardRef = /* tsx */ `
+export default React.forwardRef<HTMLAnchorElement, InternalLinkProps>(
+  function InternalLink({ variant, ...props }, ref) {
+    return (
+      <Link data-component="InternalLink"
+        ref={ref}
+        className={classNames(css.link, { [css.inverse]: variant === 'inverse' })}
+        {...props}
+      >
+        {props.children}
+      </Link>
+    );
+  }
+);`;
+
+const defaultReactForwardRefError = /* tsx */ `
+export default React.forwardRef<HTMLAnchorElement, InternalLinkProps>(
+  function InternalLink({ variant, ...props }, ref) {
+    return (
+      <Link
+        ref={ref}
+        className={classNames(css.link, { [css.inverse]: variant === 'inverse' })}
+        {...props}
+      >
+        {props.children}
+      </Link>
+    );
+  }
+);`;
+
+const defaultForwardRef = /* tsx */ `
+export default forwardRef<HTMLAnchorElement, InternalLinkProps>(
+  function InternalLink({ variant, ...props }, ref) {
+    return (
+      <Link data-component="InternalLink"
+        ref={ref}
+        className={classNames(css.link, { [css.inverse]: variant === 'inverse' })}
+        {...props}
+      >
+        {props.children}
+      </Link>
+    );
+  }
+);`;
+
+const defaultForwardRefError = /* tsx */ `
+export default forwardRef<HTMLAnchorElement, InternalLinkProps>(
+  function InternalLink({ variant, ...props }, ref) {
+    return (
+      <Link
+        ref={ref}
+        className={classNames(css.link, { [css.inverse]: variant === 'inverse' })}
+        {...props}
+      >
+        {props.children}
+      </Link>
+    );
+  }
+);`;
+
 const tests = {
   'data-component': {
     // Require the actual rule definition
@@ -193,6 +251,12 @@ const tests = {
         {
           code: forwardRef,
         },
+        {
+          code: defaultReactForwardRef,
+        },
+        {
+          code: defaultForwardRef,
+        },
       ],
       invalid: [
         {
@@ -239,6 +303,20 @@ const tests = {
             'InternalLink is missing the data-component attribute for the top-level element.',
           ],
         },
+        {
+          code: defaultReactForwardRefError,
+          output: defaultReactForwardRef,
+          errors: [
+            'InternalLink is missing the data-component attribute for the top-level element.',
+          ],
+        },
+        {
+          code: defaultForwardRefError,
+          output: defaultForwardRef,
+          errors: [
+            'InternalLink is missing the data-component attribute for the top-level element.',
+          ],
+        },
       ],
     },
   },