Class-based component support

index.js

@@ -36,6 +36,18 @@ function getReturnStatement(node) {
     return;
   }
 
+  if (node.type === 'ClassDeclaration') {
+    // For class-based components, find the render function, then its return statement
+    renderFunction = node.body?.body?.find(
+      (statement) =>
+        statement.type === 'MethodDefinition' &&
+        statement.key.name === 'render',
+    );
+    return renderFunction.value?.body?.body?.find(
+      (statement) => statement.type === 'ReturnStatement',
+    );
+  }
+
   return node.type === 'VariableDeclaration'
     ? node.declarations?.[0]?.init?.body?.body?.find(
         (statement) => statement.type === 'ReturnStatement',
@@ -99,7 +111,8 @@ const rules = {
             .filter(
               (child) =>
                 child.type === 'VariableDeclaration' ||
-                child.type === 'FunctionDeclaration',
+                child.type === 'FunctionDeclaration' ||
+                child.type === 'ClassDeclaration',
             )
             .filter((child) => {
               let flag = false;

test.js

@@ -9,6 +9,7 @@ const { join } = require('path');
 const singleComponent = `const temp = () => {
   <Icon data-component="temp" name="metric" size={24} />;
 };`;
+
 const singleComponentError = `const temp = () => {
   <Icon name="metric" size={24} />;
 };`;
@@ -88,6 +89,46 @@ const fragmentsWontUpdate = `const Component = () => {
   </>
 };`;
 
+const classComponent = `class Car extends React.Component {
+  render() {
+    return <h2 data-component="Car">Hi, I am a Car!</h2>;
+  }
+}`;
+
+const classComponentError = `class Car extends React.Component {
+  render() {
+    return <h2>Hi, I am a Car!</h2>;
+  }
+}`;
+
+const classComponentNestedError = `class Car extends React.Component {
+  render() {
+    const Door = () => (
+      <h1>I am a door!</h1>
+    );
+    return (
+      <div>
+        <Door />
+        <h2>Hi, I am a Car!</h2>
+      </div>
+    );
+  }
+}`;
+
+const classComponentNested = `class Car extends React.Component {
+  render() {
+    const Door = () => (
+      <h1>I am a door!</h1>
+    );
+    return (
+      <div data-component="Car">
+        <Door />
+        <h2>Hi, I am a Car!</h2>
+      </div>
+    );
+  }
+}`;
+
 const tests = {
   'data-component': {
     // Require the actual rule definition
@@ -142,6 +183,20 @@ const tests = {
             'temp is missing the data-component attribute for the top-level element.',
           ],
         },
+        {
+          code: classComponentError,
+          output: classComponent,
+          errors: [
+            'Car is missing the data-component attribute for the top-level element.',
+          ],
+        },
+        {
+          code: classComponentNestedError,
+          output: classComponentNested,
+          errors: [
+            'Car is missing the data-component attribute for the top-level element.',
+          ],
+        },
         {
           code: genericTestError,
           output: genericTest,