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

@@ -48,15 +48,29 @@ function getReturnStatement(node) {
     );
   }
 
+  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' &&
@@ -107,12 +121,21 @@ 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 === 'ClassDeclaration',
+                child.type === 'ClassDeclaration' ||
+                child.type === 'FunctionExpression' ||
+                child.type === 'ArrowFunctionExpression',
             )
             .filter((child) => {
               let flag = false;

test.js

@@ -132,8 +132,7 @@ const classComponentNested = `class Car extends React.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}
@@ -159,8 +158,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}
@@ -183,6 +181,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
@@ -234,6 +292,12 @@ const tests = {
         {
           code: forwardRef,
         },
+        {
+          code: defaultReactForwardRef,
+        },
+        {
+          code: defaultForwardRef,
+        },
       ],
       invalid: [
         {
@@ -294,6 +358,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.',
+          ],
+        },
       ],
     },
   },