fix: instrumentation react not to exit Class visitor early if encountering a non-class-method path

Author: artus9033

packages/instrumentation-react/src/index.ts

@@ -129,7 +129,7 @@ export default function ottrelitePlugin(): PluginObj {
 
         if (LOG_DEBUG) {
           console.log(
-            `[tracePlugin] Found directive "${traceDirective.value.value}" in ${locationDescription}, instrumenting component '${componentName}' lifecycle tracing.`
+            `[tracePlugin] Found directive "${traceDirective.value.value}" in ${locationDescription}, instrumenting function component '${componentName}' lifecycle tracing.`
           );
         }
 
@@ -240,14 +240,14 @@ export default function ottrelitePlugin(): PluginObj {
       Class(classPath, pluginPass) {
         for (const path of classPath.node.body.body) {
           if (!t.isClassMethod(path)) {
-            return;
+            continue;
           }
 
           const traceDirective = path.body.directives.find(({ value }) =>
             value.value.startsWith('use trace')
           );
 
-          if (!traceDirective) return;
+          if (!traceDirective) continue;
 
           // remove that directive
           path.body.directives = path.body.directives.filter(
@@ -266,6 +266,12 @@ export default function ottrelitePlugin(): PluginObj {
             locationDescription
           );
 
+          if (LOG_DEBUG) {
+            console.log(
+              `[tracePlugin] Found directive "${traceDirective.value.value}" in ${locationDescription}, instrumenting class component '${componentName}' lifecycle tracing.`
+            );
+          }
+
           (program as AugmentedProgramNode).__hasClassComponent = true;
 
           // wrap the class component inside the HOC

packages/instrumentation-react/src/utils/getLocationDescription.ts

@@ -2,10 +2,15 @@ import { NodePath, PluginPass } from '@babel/core';
 import * as t from '@babel/types';
 
 export function getLocationDescription(
-  path: NodePath | t.ClassMethod,
+  path: NodePath<t.Function> | t.ClassMethod | t.ClassPrivateMethod,
   pluginPass: PluginPass
 ) {
-  const loc =
-    path.type === 'ClassMethod' ? (path as t.ClassMethod).loc : path.node.loc;
+  let loc: t.SourceLocation | undefined | null;
+  if ('node' in path) {
+    loc = path.node.loc;
+  } else {
+    loc = path.loc;
+  }
+
   return `${pluginPass.filename}:${loc?.start.line}:${loc?.start.column} through ${loc?.end.line}:${loc?.end.column}`;
 }

packages/instrumentation-react/src/utils/parseDirective.ts

@@ -45,12 +45,7 @@ export function parseDirective(
   let isFunctionComponent = true;
 
   // II. Attempt to read displayName or name, if componentName is still undefined
-  if (
-    !componentName &&
-    (path.node.type === 'FunctionDeclaration' ||
-      path.node.type === 'FunctionExpression' ||
-      path.node.type === 'ArrowFunctionExpression')
-  ) {
+  if (!componentName && t.isFunction(path.node)) {
     // look for .displayName or .name assignments in parent scope
     const parent = path.findParent(
       (p) => p.isProgram() || p.isBlockStatement()
@@ -72,6 +67,14 @@ export function parseDirective(
     }
   }
 
+  const isClassMember =
+    (path.parent.type === 'ClassMethod' &&
+      path.parent.key.type === 'Identifier') ||
+    (path.parent.type === 'ObjectProperty' && t.isFunction(path.node));
+  if (isClassMember) {
+    isFunctionComponent = false;
+  }
+
   // III. Infer the name from context, if componentName is still undefined
   if (!componentName) {
     if ('id' in path.node && path.node.id?.name) {
@@ -89,13 +92,10 @@ export function parseDirective(
     ) {
       // object property assignment
       componentName ??= path.parent.key.name;
-    } else if (
-      path.parent.type === 'ClassMethod' &&
-      path.parent.key.type === 'Identifier'
-    ) {
+    } else if (isClassMember) {
       // class method
-      isFunctionComponent = false;
-      componentName ??= path.parent.key.name;
+      componentName ??= ((path.parent as t.ClassMethod).key as t.Identifier)
+        .name;
     } else if (path.parent.type === 'ExportDefaultDeclaration') {
       // anonymous export default
       throw new Error(

packages/instrumentation-react/tests/index.test.ts

@@ -395,6 +395,25 @@ describe('Ottrelite Babel Plugin', () => {
         expect(result).toContain(TRACE_COMPONENT_LIFECYCLE_HOC_NAME);
       });
 
+      test(`should add ${TRACE_COMPONENT_LIFECYCLE_HOOK_NAME} import when use trace directive is found in a class component that has other paths than the render method`, () => {
+        const code = `
+          class MyComponent {
+            someMethod() {
+              console.log('doing something');
+            }
+
+            render() {
+              'use trace';
+              return <div>Hello</div>;
+            }
+          }
+        `;
+
+        const result = transformCode(code);
+        expect(result).toContain(CORE_PACKAGE_NAME);
+        expect(result).toContain(TRACE_COMPONENT_LIFECYCLE_HOC_NAME);
+      });
+
       test(`should not add duplicate import if ${TRACE_COMPONENT_LIFECYCLE_HOOK_NAME} is already imported in a class component`, () => {
         const code = `
           import { ${TRACE_COMPONENT_LIFECYCLE_HOOK_NAME} } from '${CORE_PACKAGE_NAME}';