Handle Windows paths in python colour scheme

jonhoo · May 8, 2022 · 8d7bb2b · 8d7bb2b
1 parent 2b241f3
commit 8d7bb2b
Showing 1 changed file with 12 additions and 3 deletions.
diff --git a/src/flamegraph/color/palettes.rs b/src/flamegraph/color/palettes.rs
@@ -63,11 +63,16 @@ pub(super) mod python {
     use crate::flamegraph::color::BasicPalette;
 
     pub fn resolve(name: &str) -> BasicPalette {
-        if name.starts_with("native@") { // austin-specific format for native calls
+        if name.starts_with("native@") {
+            // austin-specific format for native calls
             return BasicPalette::Aqua;
-        } else if name.contains("/site-packages/") {
+        } else if name.contains("/site-packages/") || name.contains("\\site-packages\\") {
             return BasicPalette::Yellow;
-        } else if name.contains("/python") || name.starts_with("<frozen importlib") { // stdlib
+        } else if (name.contains("/python") && name.contains("/lib/"))
+            || (name.contains("\\Python") && name.contains("\\lib\\"))
+            || name.starts_with("<frozen importlib")
+        {
+            // stdlib
             return BasicPalette::Green;
         }
         BasicPalette::Red
@@ -465,6 +470,10 @@ mod tests {
                 input: String::from(".venv/lib/python3.9/time.py:12"),
                 output: BasicPalette::Green,
             },
+            TestData {
+                input: String::from("C:\\Users\\User\\AppData\\Local\\Programs\\Python\\Python39\\lib\\concurrent\\futures\\thread.py"),
+                output: BasicPalette::Green,
+            },
             TestData {
                 input: String::from("my_file.py:55"),
                 output: BasicPalette::Red,