Add colour palette for python traces

src/bin/flamegraph.rs

@@ -82,7 +82,7 @@ struct Opt {
         short = 'c',
         long = "colors",
         default_value = defaults::COLORS,
-        possible_values = &["aqua","blue","green","hot","io","java","js","mem","orange","perl","purple","red","rust","wakeup","yellow"],
+        possible_values = &["aqua","blue","green","hot","io","java","js","mem","orange","perl","python","purple","red","rust","wakeup","yellow","gray"],
         value_name = "STRING"
     )]
     colors: Palette,

src/flamegraph/color/mod.rs

@@ -112,6 +112,8 @@ pub enum BasicPalette {
     Purple,
     /// A palette in which colors are chosen from a orange spectrum.
     Orange,
+    /// A palette in which colors are chosen from a gray spectrum.
+    Gray,
 }
 
 /// A semantic color palette in which different hues are used to signifiy semantic aspects of
@@ -124,6 +126,8 @@ pub enum MultiPalette {
     Js,
     /// Use Perl semantics to color frames.
     Perl,
+    /// Use Python semantics to color frames.
+    Python,
     /// Use Rust semantics to color frames.
     Rust,
     /// Equivalent to [`BasicPalette::Aqua`] with [`BackgroundColor::Blue`].
@@ -220,6 +224,7 @@ impl FromStr for Palette {
             "java" => Ok(Palette::Multi(MultiPalette::Java)),
             "js" => Ok(Palette::Multi(MultiPalette::Js)),
             "perl" => Ok(Palette::Multi(MultiPalette::Perl)),
+            "python" => Ok(Palette::Multi(MultiPalette::Python)),
             "rust" => Ok(Palette::Multi(MultiPalette::Rust)),
             "red" => Ok(Palette::Basic(BasicPalette::Red)),
             "green" => Ok(Palette::Basic(BasicPalette::Green)),
@@ -228,6 +233,7 @@ impl FromStr for Palette {
             "yellow" => Ok(Palette::Basic(BasicPalette::Yellow)),
             "purple" => Ok(Palette::Basic(BasicPalette::Purple)),
             "orange" => Ok(Palette::Basic(BasicPalette::Orange)),
+            "gray" => Ok(Palette::Basic(BasicPalette::Gray)),
             unknown => Err(format!("unknown color palette: {}", unknown)),
         }
     }
@@ -331,6 +337,7 @@ fn rgb_components_for_palette(palette: Palette, name: &str, v1: f32, v2: f32, v3
         Palette::Basic(basic) => basic,
         Palette::Multi(MultiPalette::Java) => palettes::java::resolve(name),
         Palette::Multi(MultiPalette::Perl) => palettes::perl::resolve(name),
+        Palette::Multi(MultiPalette::Python) => palettes::python::resolve(name),
         Palette::Multi(MultiPalette::Js) => palettes::js::resolve(name),
         Palette::Multi(MultiPalette::Wakeup) => palettes::wakeup::resolve(name),
         Palette::Multi(MultiPalette::Rust) => palettes::rust::resolve(name),
@@ -351,6 +358,7 @@ fn rgb_components_for_palette(palette: Palette, name: &str, v1: f32, v2: f32, v3
         }
         BasicPalette::Aqua => color!(t!(50, 60_f32, v1), t!(165, 55_f32, v1), t!(165, 55_f32, v1)),
         BasicPalette::Orange => color!(t!(190, 65_f32, v1), t!(90, 65_f32, v1), t!(0, 0_f32, v1)),
+        BasicPalette::Gray => color!(t!(140, 70_f32, v1), t!(140, 70_f32, v1), t!(140, 70_f32, v1)),
     }
 }
 

src/flamegraph/color/palettes.rs

@@ -59,6 +59,31 @@ pub(super) mod perl {
     }
 }
 
+pub(super) mod python {
+    use crate::flamegraph::color::BasicPalette;
+
+    fn split_any_path(path: &str) -> impl Iterator<Item = &str> {
+        path.split(|c| c == '/' || c == '\\')
+    }
+
+    pub(in super::super) fn resolve(name: &str) -> BasicPalette {
+      if split_any_path(name).any(|part| part == "site-packages") {
+            return BasicPalette::Mem;
+        } else if split_any_path(name).any(|part| {
+            part.strip_prefix("python")
+                .or_else(|| part.strip_prefix("Python"))
+                .map_or(false, |version| {
+                    version.chars().all(|c| c.is_digit(10) || c == '.')
+                })
+        }) || name.starts_with("<")
+        {
+            // stdlib
+            return BasicPalette::Gray
+        }
+        BasicPalette::Hot
+    }
+}
+
 pub(super) mod js {
     use crate::flamegraph::color::BasicPalette;
 
@@ -104,9 +129,9 @@ pub(super) mod rust {
         if name.starts_with("core::")
             || name.starts_with("std::")
             || name.starts_with("alloc::")
-            || (name.starts_with("<core::") 
-                // Rust user-defined async functions are desugared into 
-                // GenFutures so we don't want to include those as Rust 
+            || (name.starts_with("<core::")
+                // Rust user-defined async functions are desugared into
+                // GenFutures so we don't want to include those as Rust
                 // system functions
                 && !name.starts_with("<core::future::from_generator::GenFuture<T>"))
             || name.starts_with("<std::")
@@ -437,6 +462,59 @@ mod tests {
         }
     }
 
+    #[test]
+    fn python_returns_correct() {
+        use super::python::resolve;
+
+        let test_names = [
+            TestData {
+                input: String::from("<frozen importlib._bootstrap>:_load_unlocked:680"),
+                output: BasicPalette::Green,
+            },
+            TestData {
+                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("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,
+            },
+            TestData {
+                input: String::from(".venv/lib/python3.9/site-packages/package/file.py:12"),
+                output: BasicPalette::Yellow,
+            },
+            TestData {
+                input: String::from("native@511863:_PyFunction_Vectorcall:339"),
+                output: BasicPalette::Aqua,
+            },
+            TestData {
+                input: String::from("L339"),
+                output: BasicPalette::Gray,
+            },
+            TestData {
+                input: String::from("T339"),
+                output: BasicPalette::Gray,
+            },
+            TestData {
+                input: String::from("P339"),
+                output: BasicPalette::Gray,
+            },
+        ];
+
+        for item in test_names.iter() {
+            let resolved_color = resolve(&item.input);
+            assert_eq!(resolved_color, item.output)
+        }
+    }
+
     #[test]
     fn js_returns_correct() {
         use super::js;