From 8b91cb0ef76e57526861048131e17f30748818db Mon Sep 17 00:00:00 2001
From: tgerdes <tgerdes@nvidia.com>
Date: Fri, 24 Jul 2020 11:08:58 -0500
Subject: [PATCH] Capture class information for special class functions

---
 pyprof/nvtx/nvmarker.py | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/pyprof/nvtx/nvmarker.py b/pyprof/nvtx/nvmarker.py
index be0771b..7b1baef 100644
--- a/pyprof/nvtx/nvmarker.py
+++ b/pyprof/nvtx/nvmarker.py
@@ -183,6 +183,28 @@ def get_trace_info(op_name):
             #
             fn_name = frame.name
 
+            # Capture class name for any special class functions
+            #
+            if (frame.name.startswith("__")):
+
+                # Iterate through the stack frames (like a linked list) until we get
+                # to the detailed frame we want. This is much faster and less
+                # expensive than extracting the entire frame stack every time
+                #
+                # ins stack is backwards from traceback, so depth is inverse 
+                # of current traceback depth
+                #
+                depth = len(stack) - i
+                ins_frame = ins.currentframe()
+                for _ in range(1,depth):
+                    ins_frame = ins_frame.f_back
+
+                # Grab the class name if it exists
+                #
+                if 'self' in ins_frame.f_locals:
+                    fn_name = ins_frame.f_locals['self'].__class__.__name__ + "::" + fn_name
+
+
             key = (func_stack, frame.name)
 
             if key not in func_map.keys():