TraceConfig updates to fix logic to more closely match Python `should…

…_trace` (#759) * Removed additional (redundant and outdated) checks around TraceConfig::should_trace call. * Corrected logic when profile_all is specified.
plasma-umass · Jan 18, 2024 · 5eb9a10 · 5eb9a10
1 parent bd8ea45
commit 5eb9a10
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 25 deletions.
diff --git a/src/include/traceconfig.hpp b/src/include/traceconfig.hpp
@@ -30,6 +30,11 @@ class TraceConfig {
   }
 
   bool should_trace(char* filename) {
+    if (!filename) {
+      // Defensive programming.
+      return false;
+    }
+
     auto res = _memoize.find(filename);
     if ( res != _memoize.end()) {
       return res->second;
@@ -46,19 +51,22 @@ class TraceConfig {
     const auto PATH_SEP = "/";
 #endif
 
-    auto python_lib = std::string("lib") + std::string(PATH_SEP) + std::string("python");
-    auto scalene_lib = std::string("scalene") + std::string(PATH_SEP) + std::string("scalene");
-    auto anaconda_lib = std::string("anaconda3") + std::string(PATH_SEP) + std::string("lib");
-
-    if (strstr(filename, python_lib.c_str()) != nullptr ||
-        strstr(filename, scalene_lib.c_str()) != nullptr ||
-        strstr(filename, anaconda_lib.c_str()) != nullptr ||
-	//        strstr(filename, "site-packages") != nullptr ||
-        (*filename == '<' && strstr(filename, "<ipython") != nullptr)) {
-      _memoize.insert(std::pair<std::string, bool>(std::string(filename), false));
-      return false;
+    if (!profile_all) {
+
+      auto python_lib = std::string("lib") + std::string(PATH_SEP) + std::string("python");
+      auto scalene_lib = std::string("scalene") + std::string(PATH_SEP) + std::string("scalene");
+      auto anaconda_lib = std::string("anaconda3") + std::string(PATH_SEP) + std::string("lib");
+
+      if (strstr(filename, python_lib.c_str()) ||
+	  strstr(filename, scalene_lib.c_str()) ||
+	  strstr(filename, anaconda_lib.c_str()) ||
+	  //        strstr(filename, "site-packages") != nullptr ||
+	  (strstr(filename, "<") && (strstr(filename, "<ipython") || strstr(filename, "<frozen")))) {
+	_memoize.insert(std::pair<std::string, bool>(std::string(filename), false));
+	return false;
+      }
     }
-
+    
     if (owner != nullptr) {
       for (char* traceable : items) {
         if (strstr(filename, traceable)) {

diff --git a/src/source/pywhere.cpp b/src/source/pywhere.cpp
@@ -170,6 +170,7 @@ int whereInPython(std::string& filename, int& lineno, int& bytei) {
         PyFrame_GetCode(static_cast<PyFrameObject*>(frame));
     PyPtr<> co_filename =
         PyUnicode_AsASCIIString(static_cast<PyCodeObject*>(code)->co_filename);
+
     if (!(static_cast<PyObject*>(co_filename))) {
       return 0;
     }
@@ -179,22 +180,20 @@ int whereInPython(std::string& filename, int& lineno, int& bytei) {
       continue;
     }
 
-    if (!strstr(filenameStr, "<") &&  !strstr(filenameStr, "/python") &&
-        !strstr(filenameStr, "scalene/scalene")) {
-      if (traceConfig->should_trace(filenameStr)) {
+    if (traceConfig->should_trace(filenameStr)) {
+
 #if defined(PyPy_FatalError)
-        // If this macro is defined, we are compiling PyPy, which
-        // AFAICT does not have any way to access bytecode index, so
-        // we punt and set it to 0.
-        bytei = 0;
+      // If this macro is defined, we are compiling PyPy, which
+      // AFAICT does not have any way to access bytecode index, so
+      // we punt and set it to 0.
+      bytei = 0;
 #else
-        bytei = PyFrame_GetLasti(static_cast<PyFrameObject*>(frame));
+      bytei = PyFrame_GetLasti(static_cast<PyFrameObject*>(frame));
 #endif
-        lineno = PyFrame_GetLineNumber(static_cast<PyFrameObject*>(frame));
-
-        filename = filenameStr;
-        return 1;
-      }
+      lineno = PyFrame_GetLineNumber(static_cast<PyFrameObject*>(frame));
+
+      filename = filenameStr;
+      return 1;
     }
 
     frame = PyFrame_GetBack(static_cast<PyFrameObject*>(frame));