fix: Handle module.exports = class Foo when locating classes

Author: bizob2828

src/instrumentation.rs

@@ -7,7 +7,7 @@ use std::path::PathBuf;
 use swc_core::common::{Span, SyntaxContext};
 use swc_core::ecma::{
     ast::{
-        ArrowExpr, AssignExpr, AssignTarget, BlockStmt, ClassDecl, ClassMethod, Constructor, Expr,
+        ArrowExpr, AssignExpr, AssignTarget, BlockStmt, ClassDecl, ClassExpr, ClassMethod, Constructor, Expr,
         FnDecl, FnExpr, Ident, Lit, MemberProp, MethodProp, Module, ModuleItem, Pat, PropName,
         Script, SimpleAssignTarget, Stmt, Str, VarDecl,
     },
@@ -236,6 +236,19 @@ impl Instrumentation {
             .is_none_or(|class| node.ident.sym.as_ref() == class);
         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()
+                .map_or(false, |ident| ident.sym.as_ref() == class)
+            });
+        true
+    }
 
     pub fn visit_mut_class_method(&mut self, node: &mut ClassMethod) -> bool {
         let name = match &node.key {

src/lib.rs

@@ -27,7 +27,7 @@ use swc_core::{
     common::{comments::Comments, errors::ColorConfig, FileName, FilePathMapping},
     ecma::{
         ast::{
-            AssignExpr, ClassDecl, ClassMethod, Constructor, EsVersion, FnDecl, MethodProp, Module,
+            AssignExpr, ClassDecl, ClassMethod, ClassExpr, 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_method_cjs/mod.js

@@ -7,10 +7,9 @@ class UndiciBase {
         return 42;
     }
 }
-class Undici extends UndiciBase {
+module.exports = class Undici extends UndiciBase {
     async fetch (url) {
         return super.fetch(url);
     }
-}
+};
 
-module.exports = Undici;

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'));