fix: Handle module.exports = class Foo when locating classes (#7)

Author: bizob2828

src/instrumentation.rs

@@ -7,9 +7,9 @@ use std::path::PathBuf;
 use swc_core::common::{Span, SyntaxContext};
 use swc_core::ecma::{
     ast::{
-        ArrowExpr, AssignExpr, AssignTarget, BlockStmt, ClassDecl, ClassMethod, Constructor, Expr,
-        FnDecl, FnExpr, Ident, Lit, MemberProp, MethodProp, Module, ModuleItem, Pat, PropName,
-        Script, SimpleAssignTarget, Stmt, Str, VarDecl,
+        ArrowExpr, AssignExpr, AssignTarget, BlockStmt, ClassDecl, ClassExpr, ClassMethod,
+        Constructor, Expr, FnDecl, FnExpr, Ident, Lit, MemberProp, MethodProp, Module, ModuleItem,
+        Pat, PropName, Script, SimpleAssignTarget, Stmt, Str, VarDecl,
     },
     atoms::Atom,
 };
@@ -237,6 +237,15 @@ impl Instrumentation {
         true
     }
 
+    pub fn visit_mut_class_expr(&mut self, node: &mut ClassExpr) -> bool {
+        self.is_correct_class = self.config.function_query.class_name().is_none_or(|class| {
+            node.ident
+                .as_ref()
+                .is_some_and(|ident| ident.sym.as_ref() == class)
+        });
+        true
+    }
+
     pub fn visit_mut_class_method(&mut self, node: &mut ClassMethod) -> bool {
         let name = match &node.key {
             PropName::Ident(ident) => ident.sym.clone(),

src/lib.rs

@@ -27,8 +27,8 @@ use swc_core::{
     common::{comments::Comments, errors::ColorConfig, FileName, FilePathMapping},
     ecma::{
         ast::{
-            AssignExpr, ClassDecl, ClassMethod, Constructor, EsVersion, FnDecl, MethodProp, Module,
-            Script, Str, VarDecl,
+            AssignExpr, ClassDecl, ClassExpr, ClassMethod, Constructor, EsVersion, FnDecl,
+            MethodProp, Module, Script, Str, VarDecl,
         },
         visit::{VisitMut, VisitMutWith},
     },
@@ -227,6 +227,7 @@ impl VisitMut for InstrumentationVisitor {
     visit_with_all_fn!(visit_mut_method_prop, MethodProp);
     visit_with_all_fn!(visit_mut_assign_expr, AssignExpr);
     visit_with_all_fn!(visit_mut_class_decl, ClassDecl);
+    visit_with_all_fn!(visit_mut_class_expr, ClassExpr);
     visit_with_all_fn!(visit_mut_class_method, ClassMethod);
     visit_with_all_fn!(visit_mut_constructor, Constructor);
 }

tests/class_expression_cjs/mod.js

@@ -0,0 +1,17 @@
+/**
+ * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License.
+ * This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2025 Datadog, Inc.
+ **/
+class UndiciBase {
+    async fetch (url) {
+        return 42;
+    }
+}
+class Undici extends UndiciBase {
+    async fetch (url) {
+        return super.fetch(url);
+    }
+}
+
+module.exports = Undici;
+

tests/class_expression_cjs/mod.rs

@@ -0,0 +1,15 @@
+use crate::common::*;
+use orchestrion_js::*;
+
+#[test]
+fn class_method_cjs() {
+    transpile_and_test(
+        file!(),
+        false,
+        Config::new_single(InstrumentationConfig::new(
+            "Undici:fetch",
+            test_module_matcher(),
+            FunctionQuery::class_method("Undici", "fetch", FunctionKind::Async),
+        )),
+    );
+}

tests/class_expression_cjs/test.js

@@ -0,0 +1,18 @@
+/**
+ * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License.
+ * This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2025 Datadog, Inc.
+ **/
+const Undici = require('./instrumented.js');
+const { assert, getContext } = require('../common/preamble.js');
+const context = getContext('orchestrion:undici:Undici:fetch');
+(async () => {
+  const undici = new Undici;
+  const result = await undici.fetch('https://example.com');
+  assert.strictEqual(result, 42);
+  assert.deepStrictEqual(context, {
+    start: true,
+    end: true,
+    asyncStart: 42,
+    asyncEnd: 42
+  });
+})();

tests/instrumentor_test.rs

@@ -4,6 +4,7 @@
  **/
 mod common;
 
+mod class_expression_cjs;
 mod class_method_cjs;
 mod constructor_cjs;
 mod constructor_mjs;

tests/wasm/testdata/expected-cjs.js

@@ -0,0 +1,40 @@
+const { tracingChannel: tr_ch_apm_tracingChannel } = require("diagnostics_channel");
+const tr_ch_apm$up_fetch = tr_ch_apm_tracingChannel("orchestrion:one:up:fetch");
+const tr_ch_apm$up_constructor = tr_ch_apm_tracingChannel("orchestrion:one:up:constructor");
+module.exports = class Up {
+    constructor(){
+        const tr_ch_apm_ctx$up_constructor = {
+            arguments
+        };
+        try {
+            if (tr_ch_apm$up_constructor.hasSubscribers) {
+                tr_ch_apm$up_constructor.start.publish(tr_ch_apm_ctx$up_constructor);
+            }
+            console.log('constructor');
+        } catch (tr_ch_err) {
+            if (tr_ch_apm$up_constructor.hasSubscribers) {
+                tr_ch_apm_ctx$up_constructor.error = tr_ch_err;
+                try {
+                    tr_ch_apm_ctx$up_constructor.self = this;
+                } catch (refErr) {}
+                tr_ch_apm$up_constructor.error.publish(tr_ch_apm_ctx$up_constructor);
+            }
+            throw tr_ch_err;
+        } finally{
+            if (tr_ch_apm$up_constructor.hasSubscribers) {
+                tr_ch_apm_ctx$up_constructor.self = this;
+                tr_ch_apm$up_constructor.end.publish(tr_ch_apm_ctx$up_constructor);
+            }
+        }
+    }
+    fetch() {
+        const traced = ()=>{
+            console.log('fetch');
+        };
+        if (!tr_ch_apm$up_fetch.hasSubscribers) return traced();
+        return tr_ch_apm$up_fetch.traceSync(traced, {
+            arguments,
+            self: this
+        });
+    }
+};

tests/wasm/testdata/original-cjs.js

@@ -0,0 +1,10 @@
+module.exports = class Up {
+	constructor() {
+		console.log('constructor')
+	}
+
+	fetch() {
+		console.log('fetch')
+	}
+}
+

tests/wasm/tests.mjs

@@ -33,3 +33,10 @@ const output = matchedTransforms.transform(original.toString('utf8'), true);
 
 const expected = await fs.readFile(path.join(import.meta.dirname, './testdata/expected.mjs'))
 assert.strictEqual(output, expected.toString('utf8'));
+
+const originalCjs = await fs.readFile(path.join(import.meta.dirname, './testdata/original-cjs.js'))
+const outputCjs = matchedTransforms.transform(originalCjs.toString('utf8'), false);
+
+
+const expectedCjs = await fs.readFile(path.join(import.meta.dirname, './testdata/expected-cjs.js'))
+assert.strictEqual(outputCjs, expectedCjs.toString('utf8'));