feat: wrap starlark::eval::ProfileMode

src/eval.rs

@@ -7,7 +7,7 @@ use pyo3::prelude::*;
 use pyo3::types::{PyDict, PyTuple};
 use starlark::environment::{FrozenModule, Module};
 use starlark::errors::Frame;
-use starlark::eval::{CallStack, Evaluator, FileLoader};
+use starlark::eval::{CallStack, Evaluator, FileLoader, ProfileMode};
 use starlark::PrintHandler;
 
 use crate::codemap::PyFileSpan;
@@ -252,6 +252,79 @@ impl PyEvaluator {
     }
 }
 
+/// How to profile starlark code.
+#[pyclass(
+    module = "xingque",
+    name = "ProfileMode",
+    rename_all = "SCREAMING_SNAKE_CASE",
+    frozen,
+    eq,
+    hash
+)]
+#[derive(Debug, PartialEq, Eq, Hash, Clone)]
+pub(crate) enum PyProfileMode {
+    /// The heap profile mode provides information about the time spent in each function and allocations
+    /// performed by each function. Enabling this mode the side effect of disabling garbage-collection.
+    /// This profiling mode is the recommended one.
+    HeapSummaryAllocated,
+    /// Like heap summary, but information about retained memory after module is frozen.
+    HeapSummaryRetained,
+    /// Like heap profile, but writes output comparible with
+    /// [flamegraph.pl](https://github.com/brendangregg/FlameGraph/blob/master/flamegraph.pl).
+    HeapFlameAllocated,
+    /// Like heap flame, but information about retained memory after module is frozen.
+    HeapFlameRetained,
+    /// The statement profile mode provides information about time spent in each statement.
+    Statement,
+    /// Code coverage.
+    Coverage,
+    /// The bytecode profile mode provides information about bytecode instructions.
+    Bytecode,
+    /// The bytecode profile mode provides information about bytecode instruction pairs.
+    BytecodePairs,
+    /// Provide output compatible with
+    /// [flamegraph.pl](https://github.com/brendangregg/FlameGraph/blob/master/flamegraph.pl).
+    TimeFlame,
+    /// Profile runtime typechecking.
+    Typecheck,
+}
+
+impl From<ProfileMode> for PyProfileMode {
+    fn from(value: ProfileMode) -> Self {
+        match value {
+            ProfileMode::HeapSummaryAllocated => Self::HeapSummaryAllocated,
+            ProfileMode::HeapSummaryRetained => Self::HeapSummaryRetained,
+            ProfileMode::HeapFlameAllocated => Self::HeapFlameAllocated,
+            ProfileMode::HeapFlameRetained => Self::HeapFlameRetained,
+            ProfileMode::Statement => Self::Statement,
+            ProfileMode::Coverage => Self::Coverage,
+            ProfileMode::Bytecode => Self::Bytecode,
+            ProfileMode::BytecodePairs => Self::BytecodePairs,
+            ProfileMode::TimeFlame => Self::TimeFlame,
+            ProfileMode::Typecheck => Self::Typecheck,
+            // NOTE: check if variants are added after every starlark dep bump!
+            _ => unreachable!(),
+        }
+    }
+}
+
+impl From<PyProfileMode> for ProfileMode {
+    fn from(value: PyProfileMode) -> Self {
+        match value {
+            PyProfileMode::HeapSummaryAllocated => Self::HeapSummaryAllocated,
+            PyProfileMode::HeapSummaryRetained => Self::HeapSummaryRetained,
+            PyProfileMode::HeapFlameAllocated => Self::HeapFlameAllocated,
+            PyProfileMode::HeapFlameRetained => Self::HeapFlameRetained,
+            PyProfileMode::Statement => Self::Statement,
+            PyProfileMode::Coverage => Self::Coverage,
+            PyProfileMode::Bytecode => Self::Bytecode,
+            PyProfileMode::BytecodePairs => Self::BytecodePairs,
+            PyProfileMode::TimeFlame => Self::TimeFlame,
+            PyProfileMode::Typecheck => Self::Typecheck,
+        }
+    }
+}
+
 // it would be good if https://github.com/PyO3/pyo3/issues/1190 is implemented
 // so we could have stronger typing
 // but currently duck-typing isn't bad anyway

src/lib.rs

@@ -49,6 +49,8 @@ mod xingque {
     #[pymodule_export]
     use eval::PyEvaluator;
     #[pymodule_export]
+    use eval::PyProfileMode;
+    #[pymodule_export]
     use syntax::PyAstModule;
     #[pymodule_export]
     use syntax::PyDialect;

xingque.pyi

@@ -255,6 +255,46 @@ class Evaluator:
         **kwargs: object,
     ) -> object: ...
 
+class ProfileMode:
+    """How to profile starlark code."""
+
+    HEAP_SUMMARY_ALLOCATED: ProfileMode
+    """The heap profile mode provides information about the time spent in each function and allocations
+    performed by each function. Enabling this mode the side effect of disabling garbage-collection.
+    This profiling mode is the recommended one.
+    """
+
+    HEAP_SUMMARY_RETAINED: ProfileMode
+    """Like heap summary, but information about retained memory after module is frozen."""
+
+    HEAP_FLAME_ALLOCATED: ProfileMode
+    """Like heap profile, but writes output comparible with
+    [flamegraph.pl](https://github.com/brendangregg/FlameGraph/blob/master/flamegraph.pl).
+    """
+
+    HEAP_FLAME_RETAINED: ProfileMode
+    """Like heap flame, but information about retained memory after module is frozen."""
+
+    STATEMENT: ProfileMode
+    """The statement profile mode provides information about time spent in each statement."""
+
+    COVERAGE: ProfileMode
+    """Code coverage."""
+
+    BYTECODE: ProfileMode
+    """The bytecode profile mode provides information about bytecode instructions."""
+
+    BYTECODE_PAIRS: ProfileMode
+    """The bytecode profile mode provides information about bytecode instruction pairs."""
+
+    TIME_FLAME: ProfileMode
+    """Provide output compatible with
+    [flamegraph.pl](https://github.com/brendangregg/FlameGraph/blob/master/flamegraph.pl).
+    """
+
+    TYPECHECK: ProfileMode
+    """Profile runtime typechecking."""
+
 # starlark::syntax
 
 class DialectTypes: